From 04eeec77d249b33ced9412549f672b620efd34b4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 20 May 2026 10:08:25 +0300 Subject: [PATCH] per-doc lazy items/toc cache; index document_items for fast per-doc lookup --- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 112 +++++++++--------- .../haiku/rag/store/upgrades/v0_48_0.py | 22 +++- tests/sandbox/test_sandbox_toc.py | 17 ++- 3 files changed, 82 insertions(+), 69 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index f871bb78..72b1a959 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -110,8 +110,9 @@ class Sandbox: _config: AppConfig _context: AnalysisContext _search_results: "list[SearchResult]" - _items_cache: dict[str, str] | None - _toc_cache: dict[str, str] | None + _doc_items: dict[str, list["DocumentItem"]] + _items_jsonl_cache: dict[str, str] + _toc_json_cache: dict[str, str] _repl: MontyRepl | None _vfs: OSAccess | None @@ -125,8 +126,9 @@ class Sandbox: self._config = config self._context = context self._search_results = [] - self._items_cache = None - self._toc_cache = None + self._doc_items = {} + self._items_jsonl_cache = {} + self._toc_json_cache = {} self._repl = None self._vfs = None @@ -211,40 +213,46 @@ class Sandbox: async with HaikuRAG(db_path, config=config, read_only=True) as rag: docs = await rag.list_documents(filter=self._context.filter) - doc_ids = [doc.id for doc in docs if doc.id] doc_titles = {doc.id: doc.title for doc in docs if doc.id} - def _load_caches() -> tuple[dict[str, str], dict[str, str]]: - """Bulk-fetch document items + chunk index once and build both - items.jsonl and toc.json views. Returns ``(items_cache, toc_cache)``. - """ + sandbox = self - async def _fetch() -> tuple[ - dict[str, list[DocumentItem]], dict[str, dict[str, list[str]]] - ]: + def _get_items(did: str) -> list[DocumentItem]: + """Fetch items for one doc, cached on the sandbox.""" + cached = sandbox._doc_items.get(did) + if cached is not None: + return cached + + async def _fetch() -> list[DocumentItem]: from haiku.rag.client import HaikuRAG async with HaikuRAG(db_path, config=config, read_only=True) as rag: - items_grouped = ( - await rag.document_item_repository.get_all_items_grouped( - doc_ids + return await rag.document_item_repository.get_all_items(did) + + items = _run_async(_fetch()) + sandbox._doc_items[did] = items + return items + + def _make_items_reader( + did: str, + ) -> Callable[["PurePosixPath"], str]: + def read_items(_path: "PurePosixPath") -> str: + cached = sandbox._items_jsonl_cache.get(did) + if cached is not None: + return cached + items = _get_items(did) + + async def _fetch_chunks() -> dict[str, list[str]]: + from haiku.rag.client import HaikuRAG + + async with HaikuRAG(db_path, config=config, read_only=True) as rag: + index = await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped( + [did] ) - ) - chunk_index = ( - await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped( - doc_ids - ) - ) - return items_grouped, chunk_index + return index.get(did, {}) - grouped, chunk_index = _run_async(_fetch()) - - items_cache: dict[str, str] = {} - toc_cache: dict[str, str] = {} - for did, items in grouped.items(): - doc_chunk_index = chunk_index.get(did, {}) - - items_cache[did] = "\n".join( + doc_chunk_index = _run_async(_fetch_chunks()) + jsonl = "\n".join( json.dumps( { "self_ref": item.self_ref, @@ -258,31 +266,8 @@ class Sandbox: ) for item in items ) - toc_cache[did] = json.dumps( - { - "doc_id": did, - "title": doc_titles.get(did), - "tree": _build_toc(items), - }, - ensure_ascii=False, - ) - return items_cache, toc_cache - - sandbox = self - - def _ensure_caches() -> None: - if sandbox._items_cache is None or sandbox._toc_cache is None: - items_c, toc_c = _load_caches() - sandbox._items_cache = items_c - sandbox._toc_cache = toc_c - - def _make_items_reader( - did: str, - ) -> Callable[["PurePosixPath"], str]: - def read_items(_path: "PurePosixPath") -> str: - _ensure_caches() - assert sandbox._items_cache is not None - return sandbox._items_cache.get(did, "") + sandbox._items_jsonl_cache[did] = jsonl + return jsonl return read_items @@ -290,9 +275,20 @@ class Sandbox: did: str, ) -> Callable[["PurePosixPath"], str]: def read_toc(_path: "PurePosixPath") -> str: - _ensure_caches() - assert sandbox._toc_cache is not None - return sandbox._toc_cache.get(did, "") + cached = sandbox._toc_json_cache.get(did) + if cached is not None: + return cached + items = _get_items(did) + toc = json.dumps( + { + "doc_id": did, + "title": doc_titles.get(did), + "tree": _build_toc(items), + }, + ensure_ascii=False, + ) + sandbox._toc_json_cache[did] = toc + return toc return read_toc diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py index 402d840b..c5df800f 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py @@ -1,6 +1,7 @@ import logging import pyarrow as pa +from lancedb.index import BTree from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.upgrades import Upgrade @@ -27,6 +28,21 @@ async def _ensure_columns(store: Store) -> None: ) # pragma: no cover +async def _ensure_indexes(store: Store) -> None: + """Ensure BTree scalar indexes exist on the hot document_items lookup columns. + + Fresh DBs created via ``_init_tables`` get these on first creation, but + DBs that predate that code (or were downloaded as pre-built artifacts) were + table-scanning every per-doc query — visible as ~100–300 ms even for small + docs. Built before the heading_level backfill so the per-doc WHERE clauses + in the backfill loop benefit from it. + """ + for column in ("document_id", "position", "self_ref"): + await store.document_items_table.create_index( + column, config=BTree(), replace=True + ) + + async def _apply_backfill_heading_hierarchy(store: Store) -> None: """Add heading_level + tree_depth columns and backfill from docling structure. @@ -43,6 +59,7 @@ async def _apply_backfill_heading_hierarchy(store: Store) -> None: from haiku.rag.store.models.document_item import extract_items await _ensure_columns(store) + await _ensure_indexes(store) ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist() ids = [row["id"] for row in ids] @@ -147,5 +164,8 @@ async def _apply_backfill_heading_hierarchy(store: Store) -> None: upgrade_backfill_heading_hierarchy = Upgrade( version="0.48.0", apply=_apply_backfill_heading_hierarchy, - description="Backfill heading_level + tree_depth on document_items", + description=( + "Backfill heading_level + tree_depth on document_items, " + "ensure BTree indexes on document_id / position / self_ref" + ), ) diff --git a/tests/sandbox/test_sandbox_toc.py b/tests/sandbox/test_sandbox_toc.py index a9b0db5f..494e8b6f 100644 --- a/tests/sandbox/test_sandbox_toc.py +++ b/tests/sandbox/test_sandbox_toc.py @@ -189,9 +189,9 @@ class TestTocShape: @pytest.mark.asyncio class TestTocCaching: - """toc.json is cached across reads — the items query runs once per sandbox.""" + """items + toc reads for a doc share one items fetch; repeat reads hit cache.""" - async def test_cached_across_reads(self, temp_db_path, monkeypatch): + async def test_items_fetched_once_per_doc(self, temp_db_path, monkeypatch): async with HaikuRAG(temp_db_path, create=True) as client: doc_id = await _empty_doc(client, uri="test://cache", title="Cache") await client.document_item_repository.create_items( @@ -201,21 +201,18 @@ class TestTocCaching: sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) - # Patch the repository call that backs both items.jsonl and toc.json so we - # can count how many bulk fetches happen. The sandbox opens its own - # HaikuRAG client(s) lazily; patch the class method. from haiku.rag.store.repositories.document_item import DocumentItemRepository call_count = {"n": 0} - original = DocumentItemRepository.get_all_items_grouped + original = DocumentItemRepository.get_all_items - async def counting(self, document_ids=None): + async def counting(self, document_id): call_count["n"] += 1 - return await original(self, document_ids) + return await original(self, document_id) - monkeypatch.setattr(DocumentItemRepository, "get_all_items_grouped", counting) + monkeypatch.setattr(DocumentItemRepository, "get_all_items", counting) - # First read of either file triggers ONE bulk fetch. + # Items + toc share `_doc_items`. Four reads → one items fetch. _ = await _read_toc(sandbox, doc_id) _ = await _read_items_jsonl(sandbox, doc_id) _ = await _read_toc(sandbox, doc_id)