Batch context expansion across documents

`expand_with_items` fetched its own inputs per document: one query to resolve
refs to positions, one for the window of items around them. A result set spanning
N documents cost 2N queries, which was 10 of the 18 measured for a limit=5 search
on a remote object-store corpus.

`expand_context` now does both fetches once for every document it is expanding,
and `expand_with_items` takes the positions and items it needs. Two queries for
one document, and two for five.

Each document keeps its own inclusive window in `get_items_in_ranges`. Positions
repeat across documents, so a shared range would splice one document's items into
another's context.

CHANGELOG.md

@@ -11,7 +11,7 @@
 
 ### Changed
 
-- Search enrichment issues a fixed number of `document_items` queries regardless of how many documents a result set spans, instead of one set per document.
+- Search enrichment and context expansion each issue a fixed number of `document_items` queries regardless of how many documents a result set spans, instead of one set per document. `expand_with_items` takes the items to expand from rather than fetching them, and `DocumentItemRepository` gains `resolve_refs_grouped`, `get_items_in_ranges`, `get_pictures_grouped` and `get_caption_picture_refs_grouped`.
 - Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged.
 - `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs.
 - The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift.

haiku_rag_slim/haiku/rag/client/search.py

@@ -215,7 +215,7 @@ async def expand_context(
     chunks were created without docling metadata (e.g., custom chunks passed
     to import_document).
     """
-    from haiku.rag.context import expand_with_items
+    from haiku.rag.context import expand_with_items, window_for
 
     max_chars = client._config.search.max_context_chars
 
@@ -228,24 +228,38 @@ async def expand_context(
         document_groups[doc_id].append(result)
 
     expanded_results = []
-
+    expandable = {
+        doc_id: doc_results
+        for doc_id, doc_results in document_groups.items()
+        if doc_id is not None and any(r.doc_item_refs for r in doc_results)
+    }
     for doc_id, doc_results in document_groups.items():
-        if doc_id is None:
+        if doc_id not in expandable:
             expanded_results.extend(doc_results)
-            continue
 
-        has_refs = any(r.doc_item_refs for r in doc_results)
+    repo = client.document_item_repository
-        if not has_refs:
+    positions_by_document = await repo.resolve_refs_grouped(
-            expanded_results.extend(doc_results)
+        {
-            continue
+            doc_id: [ref for r in doc_results for ref in r.doc_item_refs]
+            for doc_id, doc_results in expandable.items()
+        }
+    )
+    windows = {
+        doc_id: window_for(positions)
+        for doc_id, positions in positions_by_document.items()
+        if positions
+    }
+    items_by_document = await repo.get_items_in_ranges(windows)
 
-        expanded = await expand_with_items(
+    for doc_id, doc_results in expandable.items():
-            client.document_item_repository,
+        expanded_results.extend(
-            doc_id,
+            expand_with_items(
-            doc_results,
+                doc_results,
-            max_chars,
+                max_chars,
+                positions_by_document.get(doc_id, {}),
+                items_by_document.get(doc_id, []),
+            )
         )
-        expanded_results.extend(expanded)
 
     expanded_results.sort(key=lambda r: r.score, reverse=True)
     # image_data and picture_captions are preserved through expansion by

haiku_rag_slim/haiku/rag/context.py

@@ -34,7 +34,6 @@ In both cases:
 
 from haiku.rag.store.models.chunk import SearchResult
 from haiku.rag.store.models.document_item import DocumentItem
-from haiku.rag.store.repositories.document_item import DocumentItemRepository
 
 _NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
 _SECTION_BOUNDARY_LABELS = {"section_header", "title"}
@@ -425,33 +424,28 @@ def _build_result(
 _WINDOW_MARGIN = 100
 
 
-async def expand_with_items(
+def window_for(ref_positions: dict[str, int]) -> tuple[int, int]:
-    document_item_repository: DocumentItemRepository,
+    """The inclusive position range to fetch around a document's matches.
-    document_id: str,
+
+    The margin must be wide enough to find section boundaries: the nearest
+    section_header or title above and below the match.
+    """
+    positions = sorted(ref_positions.values())
+    return max(0, positions[0] - _WINDOW_MARGIN), positions[-1] + _WINDOW_MARGIN
+
+
+def expand_with_items(
     results: list[SearchResult],
     max_chars: int,
+    ref_positions: dict[str, int],
+    window_items: list[DocumentItem],
 ) -> list[SearchResult]:
-    """Expand results using the document_items table."""
+    """Expand results from items already fetched.
-    all_refs = []
-    for result in results:
-        all_refs.extend(result.doc_item_refs)
 
-    ref_positions = await document_item_repository.resolve_refs(document_id, all_refs)
+    Fetching is the caller's, so one query can serve every document in a result
-    if not ref_positions:
+    set rather than one per document.
-        return results
+    """
-
+    if not ref_positions or not window_items:
-    # Fetch a window of items around matched positions. The margin must be
-    # wide enough to find section boundaries (the nearest section_header/title
-    # above and below the match).
-    all_positions = sorted(ref_positions.values())
-    window_margin = _WINDOW_MARGIN
-    window_start = max(0, min(all_positions) - window_margin)
-    window_end = max(all_positions) + window_margin
-    window_items = await document_item_repository.get_items_in_range(
-        document_id, window_start, window_end
-    )
-
-    if not window_items:
         return results
 
     has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items)

haiku_rag_slim/haiku/rag/store/repositories/document_item.py

@@ -167,6 +167,58 @@ class DocumentItemRepository:
         )
         return {row["self_ref"]: row["position"] for row in rows}
 
+    async def resolve_refs_grouped(
+        self, refs_by_document: "Mapping[str, Sequence[str]]"
+    ) -> dict[str, dict[str, int]]:
+        """`resolve_refs` across documents in one query."""
+        predicate = self._per_document_predicate(refs_by_document, "self_ref")
+        if predicate is None:
+            return {}
+        rows = await (
+            self.store.document_items_table.query()
+            .select(["document_id", "self_ref", "position"])
+            .where(predicate)
+            .to_list()
+        )
+        grouped: dict[str, dict[str, int]] = {}
+        for row in rows:
+            grouped.setdefault(row["document_id"], {})[row["self_ref"]] = row[
+                "position"
+            ]
+        return grouped
+
+    async def get_items_in_ranges(
+        self, ranges_by_document: "Mapping[str, tuple[int, int]]"
+    ) -> dict[str, list[DocumentItem]]:
+        """`get_items_in_range` across documents in one query.
+
+        Each document keeps its own inclusive range. Positions repeat across
+        documents, so a shared range would splice one document's items into
+        another's context.
+        """
+        clauses = []
+        for document_id, (start, end) in ranges_by_document.items():
+            safe_id = escape_sql_string(document_id)
+            clauses.append(
+                f"(document_id = '{safe_id}' "
+                f"AND position >= {start} AND position <= {end})"
+            )
+        if not clauses:
+            return {}
+        rows = await (
+            self.store.document_items_table.query()
+            .select(_METADATA_COLUMNS)
+            .where(" OR ".join(clauses))
+            .to_list()
+        )
+        grouped: dict[str, list[DocumentItem]] = {}
+        for row in rows:
+            item = self._record_to_item(row)
+            grouped.setdefault(item.document_id, []).append(item)
+        for items in grouped.values():
+            items.sort(key=lambda x: x.position)
+        return grouped
+
     async def get_item_count(self, document_id: str) -> int:
         """Count items for a document."""
         safe_id = escape_sql_string(document_id)

tests/test_context.py

@@ -8,6 +8,7 @@ from haiku.rag.context import (
     _find_expansion_range,
     _merge_ranges,
     expand_with_items,
+    window_for,
 )
 from haiku.rag.store.models.chunk import SearchResult
 from haiku.rag.store.models.document_item import DocumentItem
@@ -328,6 +329,22 @@ class TestClipToBudget:
 
 
 @pytest.mark.asyncio
+async def _fetch_and_expand(repo, document_id, results, max_chars):
+    """Do the fetching `expand_context` does, for tests exercising the
+    expansion logic rather than the batched fetch."""
+    positions = (
+        await repo.resolve_refs_grouped(
+            {document_id: [ref for r in results for ref in r.doc_item_refs]}
+        )
+    ).get(document_id, {})
+    items: list = []
+    if positions:
+        items = (
+            await repo.get_items_in_ranges({document_id: window_for(positions)})
+        ).get(document_id, [])
+    return expand_with_items(results, max_chars, positions, items)
+
+
 class TestExpandWithItems:
     async def test_unresolvable_refs_returns_original(self, temp_db_path):
         from haiku.rag.client import HaikuRAG
@@ -349,7 +366,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/999999"],
             )
             assert doc.id is not None
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, doc.id, [result], 5000
             )
             assert len(expanded) == 1
@@ -399,7 +416,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/1"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -480,7 +497,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/pictures/0"],
                 page_numbers=[13],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -542,7 +559,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/1", "#/texts/2", "#/texts/3", "#/texts/4"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -574,7 +591,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/tables/0"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 10_000
             )
             assert len(expanded) == 1
@@ -612,7 +629,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/0"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 10_000
             )
             assert len(expanded) == 1
@@ -655,7 +672,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/1"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -690,7 +707,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/0"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -720,7 +737,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/2"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -756,7 +773,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/3"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r1, r2], 5000
             )
             # Ranges around positions 1 and 3 overlap → one merged result.
@@ -791,7 +808,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/1"],
                 document_meta={"source_url": "https://example.org/report/view"},
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r1], 5000
             )
             assert len(expanded) == 1
@@ -832,7 +849,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/3"],
                 chunk_meta={"para_no": "14"},
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r1, r2], 5000
             )
             assert len(expanded) == 1
@@ -873,7 +890,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/3"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r_early, r_best], 5000
             )
             assert len(expanded) == 1
@@ -934,7 +951,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/2"],
                 page_numbers=[3],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r_low, r_high], 500
             )
             # The clip window around HIGHMARK cannot contain LOWMARK's item,
@@ -990,7 +1007,7 @@ class TestExpandWithItems:
                     page_numbers=[2],
                 ),
             ]
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", inputs, 100
             )
             assert len(expanded) == 2
@@ -1037,7 +1054,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/texts/5"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r1, r2], 400
             )
             assert len(expanded) == 1
@@ -1089,7 +1106,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/1"],
                 page_numbers=[2],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r1, r2], 5000
             )
             assert len(expanded) == 2
@@ -1143,7 +1160,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/1"],
                 page_numbers=[8],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository,
                 "doc-1",
                 [r_missing_item_page, r_with_item_page],
@@ -1190,7 +1207,7 @@ class TestExpandWithItems:
                 doc_item_refs=["#/texts/0", "#/texts/1"],
                 page_numbers=[1, 2],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 100
             )
 
@@ -1226,7 +1243,7 @@ class TestExpandWithItems:
                 document_id="doc-1",
                 doc_item_refs=["#/tables/0"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 10_000
             )
             assert len(expanded) == 1
@@ -1289,7 +1306,7 @@ class TestExpandWithItemsPictureBytes:
                 image_data={"#/pictures/0": "BASE64BYTES"},
                 picture_captions={"#/pictures/0": "Figure 1 caption."},
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [result], 5000
             )
             assert len(expanded) == 1
@@ -1328,7 +1345,7 @@ class TestExpandWithItemsPictureBytes:
                 doc_item_refs=["#/pictures/3"],
                 image_data={"#/pictures/3": "B"},
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r1, r2], 5000
             )
             # Ranges around positions 1 and 3 overlap → one merged result.
@@ -1389,7 +1406,7 @@ class TestExpandWithItemsPictureBytes:
                 doc_item_refs=["#/pictures/1"],
                 image_data={"#/pictures/1": "HIGHBYTES"},
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, "doc-1", [r_low, r_high], 500
             )
             assert len(expanded) == 2
@@ -1457,7 +1474,7 @@ class TestExpandWithItemsWindowEdges:
                 document_id=doc.id,
                 doc_item_refs=["#/texts/0"],
             )
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, doc.id, [result], 5000
             )
 
@@ -1502,7 +1519,7 @@ class TestExpandWithItemsWindowEdges:
                 doc_item_refs=["#/texts/404"],
             )
 
-            expanded = await expand_with_items(
+            expanded = await _fetch_and_expand(
                 rag.document_item_repository, doc.id, [resolvable, unmatched], 5000
             )
 

tests/test_enrichment_batching.py

@@ -132,3 +132,81 @@ async def test_caption_ranked_results_take_at_most_four_queries(
             assert all(r.image_data for r in results)
 
     assert counts == [3, 3], counts
+
+
+async def _seed_expandable(rag: HaikuRAG, document_ids: list[str]) -> None:
+    """A section header and two text items, so expansion has something to widen
+    into. Positions and self_refs repeat across documents."""
+    for document_id in document_ids:
+        await rag.document_item_repository.create_items(
+            document_id,
+            [
+                DocumentItem(
+                    document_id=document_id,
+                    position=0,
+                    self_ref="#/texts/0",
+                    label="section_header",
+                    text=f"Section of {document_id}",
+                ),
+                DocumentItem(
+                    document_id=document_id,
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="text",
+                    text=f"anchor body of {document_id}",
+                ),
+                DocumentItem(
+                    document_id=document_id,
+                    position=2,
+                    self_ref="#/texts/2",
+                    label="text",
+                    text=f"neighbouring body of {document_id}",
+                ),
+            ],
+        )
+
+
+def _text_result(document_id: str) -> SearchResult:
+    return SearchResult(
+        chunk_id=f"{document_id}-anchor",
+        document_id=document_id,
+        content=f"anchor body of {document_id}",
+        score=0.9,
+        doc_item_refs=["#/texts/1"],
+    )
+
+
+@pytest.mark.asyncio
+async def test_expansion_query_count_is_flat_in_document_count(
+    temp_db_path, item_queries
+):
+    async with HaikuRAG(temp_db_path, create=True) as rag:
+        await _seed_expandable(rag, [f"doc-{i}" for i in range(5)])
+
+        counts = []
+        for n in (1, 5):
+            results = [_text_result(f"doc-{i}") for i in range(n)]
+            item_queries["n"] = 0
+            expanded = await rag.expand_context(results)
+            counts.append(item_queries["n"])
+            assert len(expanded) == n
+
+    assert counts == [2, 2], counts
+
+
+@pytest.mark.asyncio
+async def test_expansion_widens_each_document_with_its_own_items(temp_db_path):
+    """Positions repeat across documents, so a batched window fetch keyed on
+    position alone would splice one document's text into another's context."""
+    async with HaikuRAG(temp_db_path, create=True) as rag:
+        await _seed_expandable(rag, ["doc-a", "doc-b"])
+
+        expanded = await rag.expand_context(
+            [_text_result("doc-a"), _text_result("doc-b")]
+        )
+
+    by_doc = {r.document_id: r.content for r in expanded}
+    assert "neighbouring body of doc-a" in by_doc["doc-a"]
+    assert "doc-b" not in by_doc["doc-a"]
+    assert "neighbouring body of doc-b" in by_doc["doc-b"]
+    assert "doc-a" not in by_doc["doc-b"]

tests/test_picture_in_context.py

@@ -16,10 +16,10 @@ from haiku.rag.capabilities.rag import RAGState, create_capability
 from haiku.rag.client import HaikuRAG
 from haiku.rag.client.search import _populate_image_data
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.context import expand_with_items
 from haiku.rag.store.models.chunk import SearchResult
 from haiku.rag.store.models.document_item import DocumentItem
 from haiku.rag.tools.search import create_search_toolset
+from tests.test_context import _fetch_and_expand
 
 
 def _make_png(color: str = "red", size: tuple[int, int] = (4, 4)) -> bytes:
@@ -205,11 +205,8 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat
             doc_item_refs=["#/texts/1"],
             labels=["paragraph"],
         )
-        expanded = await expand_with_items(
+        expanded = await _fetch_and_expand(
-            rag.document_item_repository,
+            rag.document_item_repository, "doc-1", [seed], 10_000
-            "doc-1",
-            [seed],
-            max_chars=10_000,
         )
         assert len(expanded) == 1
         out = expanded[0]