From 5b0444043af807ed4fcc1657c9a4459f50202be9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 16:35:17 +0300 Subject: [PATCH] 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 | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 42 ++++++---- haiku_rag_slim/haiku/rag/context.py | 42 +++++----- .../rag/store/repositories/document_item.py | 52 +++++++++++++ tests/test_context.py | 67 ++++++++++------ tests/test_enrichment_batching.py | 78 +++++++++++++++++++ tests/test_picture_in_context.py | 9 +-- 7 files changed, 222 insertions(+), 70 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 332bfb8a..4ea0ee94 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 6f84e9d8..71bae916 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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) - if not has_refs: - expanded_results.extend(doc_results) - continue + repo = client.document_item_repository + positions_by_document = await repo.resolve_refs_grouped( + { + 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( - client.document_item_repository, - doc_id, - doc_results, - max_chars, + for doc_id, doc_results in expandable.items(): + expanded_results.extend( + expand_with_items( + doc_results, + 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 diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 42073423..9aa9a486 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/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( - document_item_repository: DocumentItemRepository, - document_id: str, +def window_for(ref_positions: dict[str, int]) -> tuple[int, int]: + """The inclusive position range to fetch around a document's matches. + + 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.""" - all_refs = [] - for result in results: - all_refs.extend(result.doc_item_refs) + """Expand results from items already fetched. - ref_positions = await document_item_repository.resolve_refs(document_id, all_refs) - if not ref_positions: - return results - - # 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: + Fetching is the caller's, so one query can serve every document in a result + set rather than one per document. + """ + if not ref_positions or not window_items: return results has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 3c3f1662..7932868d 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/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) diff --git a/tests/test_context.py b/tests/test_context.py index 4d307f60..2d73c792 100644 --- a/tests/test_context.py +++ b/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 ) diff --git a/tests/test_enrichment_batching.py b/tests/test_enrichment_batching.py index 0e3a8ff5..a5aa15c3 100644 --- a/tests/test_enrichment_batching.py +++ b/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"] diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index a9b32a47..944afd56 100644 --- a/tests/test_picture_in_context.py +++ 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( - rag.document_item_repository, - "doc-1", - [seed], - max_chars=10_000, + expanded = await _fetch_and_expand( + rag.document_item_repository, "doc-1", [seed], 10_000 ) assert len(expanded) == 1 out = expanded[0]