per-doc lazy items/toc cache; index document_items for fast per-doc lookup

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
+    _doc_items: dict[str, list["DocumentItem"]]
-    _toc_cache: dict[str, str] | None
+    _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._doc_items = {}
-        self._toc_cache = None
+        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]]:
+        sandbox = self
-            """Bulk-fetch document items + chunk index once and build both
-            items.jsonl and toc.json views. Returns ``(items_cache, toc_cache)``.
-            """
 
-            async def _fetch() -> tuple[
+        def _get_items(did: str) -> list[DocumentItem]:
-                dict[str, list[DocumentItem]], dict[str, dict[str, list[str]]]
+            """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 = (
+                    return await rag.document_item_repository.get_all_items(did)
-                        await rag.document_item_repository.get_all_items_grouped(
+
-                            doc_ids
+            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]
                         )
-                    )
+                    return index.get(did, {})
-                    chunk_index = (
-                        await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
-                            doc_ids
-                        )
-                    )
-                return items_grouped, chunk_index
 
-            grouped, chunk_index = _run_async(_fetch())
+                doc_chunk_index = _run_async(_fetch_chunks())
-
+                jsonl = "\n".join(
-            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(
                     json.dumps(
                         {
                             "self_ref": item.self_ref,
@@ -258,31 +266,8 @@ class Sandbox:
                     )
                     for item in items
                 )
-                toc_cache[did] = json.dumps(
+                sandbox._items_jsonl_cache[did] = jsonl
-                    {
+                return jsonl
-                        "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, "")
 
             return read_items
 
@@ -290,9 +275,20 @@ class Sandbox:
             did: str,
         ) -> Callable[["PurePosixPath"], str]:
             def read_toc(_path: "PurePosixPath") -> str:
-                _ensure_caches()
+                cached = sandbox._toc_json_cache.get(did)
-                assert sandbox._toc_cache is not None
+                if cached is not None:
-                return sandbox._toc_cache.get(did, "")
+                    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
 

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"
+    ),
 )

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)