danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-20 10:08:25 +03:00
parent e3314e36fd
commit 04eeec77d2
No known key found for this signature in database
3 changed files with 82 additions and 69 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

112
haiku_rag_slim/haiku/rag/sandbox/sandbox.py
View file

@ -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

22
haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py
View file

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

17
tests/sandbox/test_sandbox_toc.py
View file

@ -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)