From af6a6b0bbe6ea6982bb6d4fa535541c788380863 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 15:50:42 +0300 Subject: [PATCH 1/6] Batch search enrichment across documents MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_populate_image_data` ran its stages once per result document, so a result set spanning N documents cost 4N `document_items` queries. Measured on a remote object-store corpus, a limit=5 search with expansion was 18 queries, 16 of them against `document_items`. The stages now run once each across every document, and flat in document count: two queries for the dependent caption-to-picture mapping when results ranked on a caption, one for the picture bytes. Two queries for a picture-ref result set, three at most. Picture text comes back with the bytes rather than from a second query, since it is on the same rows. Predicates are per document, `(document_id = 'a' AND self_ref IN (…)) OR (…)`, rather than `self_ref IN (union)`. self_ref and position values repeat across documents, so a union predicate would return other documents' rows: for picture_data that fetches blobs nobody asked for, and it can hand one document another document's picture. --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/client/search.py | 68 +++++---- .../rag/store/repositories/document_item.py | 100 +++++++++++++ tests/store/test_document_item_grouped.py | 116 +++++++++++++++ tests/test_enrichment_batching.py | 134 ++++++++++++++++++ 5 files changed, 388 insertions(+), 31 deletions(-) create mode 100644 tests/store/test_document_item_grouped.py create mode 100644 tests/test_enrichment_batching.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 46d8f372..332bfb8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +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. - 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 46752242..6f84e9d8 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -139,22 +139,27 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) A result carries a picture when its refs include the picture directly, or when they include the picture's caption — the common case where a prose chunk carrying a figure's caption ranks while the picture is its own chunk. - Groups results by document_id and batches one picture-bytes lookup per - document so a result set spanning N documents costs N reads, not one per - picture. + Costs a fixed number of reads however many documents the result set spans. """ repo = client.document_item_repository by_doc: dict[str, list[SearchResult]] = {} for r in results: if r.document_id and r.doc_item_refs: by_doc.setdefault(r.document_id, []).append(r) + if not by_doc: + return + refs_by_document = { + doc_id: list({ref for r in doc_results for ref in r.doc_item_refs}) + for doc_id, doc_results in by_doc.items() + } + captions_to_pictures = await repo.get_caption_picture_refs_grouped(refs_by_document) + + # Which pictures each result wants, and which to fetch per document. + result_pictures: list[tuple[SearchResult, list[str]]] = [] + wanted: dict[str, list[str]] = {} for doc_id, doc_results in by_doc.items(): - all_refs = {ref for r in doc_results for ref in r.doc_item_refs} - caption_to_picture = await repo.get_caption_picture_refs(doc_id, list(all_refs)) - - result_pictures: list[tuple[SearchResult, list[str]]] = [] - wanted: list[str] = [] + caption_to_picture = captions_to_pictures.get(doc_id, {}) seen: set[str] = set() for r in doc_results: pictures: list[str] = [] @@ -170,30 +175,31 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) result_pictures.append((r, pictures)) for picture in pictures: if picture not in seen: - wanted.append(picture) + wanted.setdefault(doc_id, []).append(picture) seen.add(picture) - if not wanted: - continue - bytes_by_ref = await repo.get_pictures_for_chunk(doc_id, wanted) - if not bytes_by_ref: - continue - captions_by_ref = await repo.get_text_for_refs( - doc_id, list(bytes_by_ref.keys()) - ) - for r, pictures in result_pictures: - attached: dict[str, str] = {} - captions: dict[str, str] = {} - for ref in pictures: - blob = bytes_by_ref.get(ref) - if blob: - attached[ref] = base64.b64encode(blob).decode("ascii") - caption = captions_by_ref.get(ref) - if caption: - captions[ref] = caption - if attached: - r.image_data = attached - if captions: - r.picture_captions = captions + if not wanted: + return + + bytes_by_document, captions_by_document = await repo.get_pictures_grouped(wanted) + if not bytes_by_document: + return + + for r, pictures in result_pictures: + bytes_by_ref = bytes_by_document.get(r.document_id or "", {}) + captions_by_ref = captions_by_document.get(r.document_id or "", {}) + attached: dict[str, str] = {} + captions: dict[str, str] = {} + for ref in pictures: + blob = bytes_by_ref.get(ref) + if blob: + attached[ref] = base64.b64encode(blob).decode("ascii") + caption = captions_by_ref.get(ref) + if caption: + captions[ref] = caption + if attached: + r.image_data = attached + if captions: + r.picture_captions = captions async def expand_context( 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 cc26226d..3c3f1662 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping, Sequence from haiku.rag.store.engine import DocumentItemRecord, Store from haiku.rag.store.models.document_item import DocumentItem @@ -242,6 +243,105 @@ class DocumentItemRepository: result[row["self_ref"]] = data return result + @staticmethod + def _per_document_predicate( + refs_by_document: "Mapping[str, Sequence[str | int]]", column: str + ) -> str | None: + """`(document_id = 'a' AND col IN (...)) OR (document_id = 'b' AND ...)`. + + Per-document rather than `col IN (union)`: self_ref and position values + repeat across documents, so a union predicate would return other + documents' rows, which for picture_data means fetching blobs nobody asked + for. Returns None when nothing is asked for. + """ + clauses = [] + for document_id, refs in refs_by_document.items(): + if not refs: + continue + safe_id = escape_sql_string(document_id) + values = ", ".join( + str(r) if isinstance(r, int) else f"'{escape_sql_string(r)}'" + for r in refs + ) + clauses.append(f"(document_id = '{safe_id}' AND {column} IN ({values}))") + return " OR ".join(clauses) if clauses else None + + async def get_pictures_grouped( + self, refs_by_document: "Mapping[str, list[str]]" + ) -> tuple[dict[str, dict[str, bytes]], dict[str, dict[str, str]]]: + """Picture bytes and their text, across documents, in one query. + + Returns `(bytes_by_document, text_by_document)`, each + `{document_id: {self_ref: value}}` and each omitting refs whose value is + empty. The text comes from the same rows as the bytes, so asking for it + separately would be a second read of rows already in hand. + """ + 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", "picture_data", "text"]) + .where(predicate) + .to_list() + ) + blobs: dict[str, dict[str, bytes]] = {} + texts: dict[str, dict[str, str]] = {} + for row in rows: + data = row.get("picture_data") + if not data: + continue + blobs.setdefault(row["document_id"], {})[row["self_ref"]] = data + text = row.get("text") + if text: + texts.setdefault(row["document_id"], {})[row["self_ref"]] = text + return blobs, texts + + async def get_caption_picture_refs_grouped( + self, refs_by_document: "Mapping[str, list[str]]" + ) -> dict[str, dict[str, str]]: + """`get_caption_picture_refs` across documents in two queries. + + Two rather than one because the stages are dependent: a caption's + picture is the item at `position - 1`, which the first query is what + establishes. + """ + predicate = self._per_document_predicate(refs_by_document, "self_ref") + if predicate is None: + return {} + caption_rows = await ( + self.store.document_items_table.query() + .select(["document_id", "self_ref", "position"]) + .where(f"label = 'caption' AND ({predicate})") + .to_list() + ) + if not caption_rows: + return {} + + prev_to_caption: dict[str, dict[int, str]] = {} + for row in caption_rows: + prev_to_caption.setdefault(row["document_id"], {})[row["position"] - 1] = ( + row["self_ref"] + ) + + # Non-empty: every caption row contributed a position. + picture_predicate = self._per_document_predicate( + {did: list(positions) for did, positions in prev_to_caption.items()}, + "position", + ) + picture_rows = await ( + self.store.document_items_table.query() + .select(["document_id", "self_ref", "position"]) + .where(f"label = 'picture' AND ({picture_predicate})") + .to_list() + ) + grouped: dict[str, dict[str, str]] = {} + for row in picture_rows: + caption = prev_to_caption.get(row["document_id"], {}).get(row["position"]) + if caption: + grouped.setdefault(row["document_id"], {})[caption] = row["self_ref"] + return grouped + async def get_text_for_refs( self, document_id: str, refs: list[str] ) -> dict[str, str]: diff --git a/tests/store/test_document_item_grouped.py b/tests/store/test_document_item_grouped.py new file mode 100644 index 00000000..2daf03c7 --- /dev/null +++ b/tests/store/test_document_item_grouped.py @@ -0,0 +1,116 @@ +import lancedb +import pytest + +from haiku.rag.store.engine import Store +from haiku.rag.store.models import DocumentItem +from haiku.rag.store.repositories.document_item import DocumentItemRepository + + +async def _seed(repo: DocumentItemRepository, document_id: str) -> None: + """A picture at position 0 with its caption at position 1, the docling + layout. Every document uses the same self_refs.""" + await repo.create_items( + document_id, + [ + DocumentItem( + document_id=document_id, + position=0, + self_ref="#/pictures/0", + label="picture", + text=f"caption text {document_id}", + picture_data=f"bytes-{document_id}".encode(), + ), + DocumentItem( + document_id=document_id, + position=1, + self_ref="#/texts/1", + label="caption", + text=f"figure 1 {document_id}", + ), + ], + ) + + +@pytest.fixture +def item_queries(monkeypatch): + tally = {"n": 0} + query = lancedb.AsyncTable.query + + def counted(self): + if self.name == "document_items": + tally["n"] += 1 + return query(self) + + monkeypatch.setattr(lancedb.AsyncTable, "query", counted) + return tally + + +@pytest.mark.asyncio +async def test_pictures_grouped_keeps_documents_apart(temp_db_path, item_queries): + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + await _seed(repo, "doc-a") + await _seed(repo, "doc-b") + + item_queries["n"] = 0 + blobs, texts = await repo.get_pictures_grouped( + {"doc-a": ["#/pictures/0"], "doc-b": ["#/pictures/0"]} + ) + + assert item_queries["n"] == 1 + assert blobs == { + "doc-a": {"#/pictures/0": b"bytes-doc-a"}, + "doc-b": {"#/pictures/0": b"bytes-doc-b"}, + } + assert texts == { + "doc-a": {"#/pictures/0": "caption text doc-a"}, + "doc-b": {"#/pictures/0": "caption text doc-b"}, + } + + +@pytest.mark.asyncio +async def test_pictures_grouped_fetches_only_requested_documents(temp_db_path): + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + await _seed(repo, "doc-a") + await _seed(repo, "doc-b") + + blobs, _ = await repo.get_pictures_grouped({"doc-a": ["#/pictures/0"]}) + + assert blobs == {"doc-a": {"#/pictures/0": b"bytes-doc-a"}} + + +@pytest.mark.asyncio +async def test_caption_picture_refs_grouped_uses_two_queries( + temp_db_path, item_queries +): + """The stages are dependent: the caption's position is what finds its + picture, so this is two queries however many documents are asked for.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + for document_id in ("doc-a", "doc-b", "doc-c"): + await _seed(repo, document_id) + + item_queries["n"] = 0 + got = await repo.get_caption_picture_refs_grouped( + {did: ["#/texts/1"] for did in ("doc-a", "doc-b", "doc-c")} + ) + + assert item_queries["n"] == 2 + for document_id in ("doc-a", "doc-b", "doc-c"): + assert got[document_id] == {"#/texts/1": "#/pictures/0"} + + +@pytest.mark.asyncio +async def test_grouped_calls_with_nothing_asked_for_do_not_query( + temp_db_path, item_queries +): + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + + item_queries["n"] = 0 + assert await repo.get_pictures_grouped({}) == ({}, {}) + assert await repo.get_pictures_grouped({"doc-a": []}) == ({}, {}) + assert await repo.get_caption_picture_refs_grouped({}) == {} + + assert item_queries["n"] == 0 diff --git a/tests/test_enrichment_batching.py b/tests/test_enrichment_batching.py new file mode 100644 index 00000000..0e3a8ff5 --- /dev/null +++ b/tests/test_enrichment_batching.py @@ -0,0 +1,134 @@ +import lancedb +import pytest + +from haiku.rag.client import HaikuRAG +from haiku.rag.client.search import _populate_image_data +from haiku.rag.store.models import DocumentItem, SearchResult + + +def _picture_result(document_id: str, ref: str) -> SearchResult: + return SearchResult( + chunk_id=f"{document_id}-{ref}", + document_id=document_id, + content="body", + score=0.9, + doc_item_refs=[ref], + ) + + +async def _seed(rag: HaikuRAG, document_ids: list[str]) -> None: + """Each document gets the same self_refs, which is what real documents do: + `#/pictures/0` exists in every one of them.""" + for document_id in document_ids: + await rag.document_item_repository.create_items( + document_id, + [ + DocumentItem( + document_id=document_id, + position=0, + self_ref="#/pictures/0", + label="picture", + text=f"caption for {document_id}", + picture_data=f"bytes-{document_id}".encode(), + ), + ], + ) + + +@pytest.fixture +def item_queries(monkeypatch): + tally = {"n": 0} + query = lancedb.AsyncTable.query + + def counted(self): + if self.name == "document_items": + tally["n"] += 1 + return query(self) + + monkeypatch.setattr(lancedb.AsyncTable, "query", counted) + return tally + + +@pytest.mark.asyncio +async def test_enrichment_query_count_does_not_grow_with_documents( + temp_db_path, item_queries +): + async with HaikuRAG(temp_db_path, create=True) as rag: + await _seed(rag, [f"doc-{i}" for i in range(6)]) + + results = [_picture_result("doc-0", "#/pictures/0")] + item_queries["n"] = 0 + await _populate_image_data(rag, results) + one_document = item_queries["n"] + + results = [_picture_result(f"doc-{i}", "#/pictures/0") for i in range(6)] + item_queries["n"] = 0 + await _populate_image_data(rag, results) + six_documents = item_queries["n"] + + assert (one_document, six_documents) == (2, 2), ( + f"one document took {one_document} queries, six took {six_documents}" + ) + + +@pytest.mark.asyncio +async def test_each_document_gets_its_own_pictures(temp_db_path): + """self_refs collide across documents, so a batched fetch keyed on self_ref + alone would hand one document another's picture.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + await _seed(rag, ["doc-a", "doc-b"]) + + results = [ + _picture_result("doc-a", "#/pictures/0"), + _picture_result("doc-b", "#/pictures/0"), + ] + await _populate_image_data(rag, results) + + import base64 + + for result, document_id in zip(results, ["doc-a", "doc-b"]): + assert result.image_data is not None + blob = base64.b64decode(result.image_data["#/pictures/0"]) + assert blob == f"bytes-{document_id}".encode() + + +@pytest.mark.asyncio +async def test_caption_ranked_results_take_at_most_four_queries( + temp_db_path, item_queries +): + """The worst case: results ranked on a caption, so the dependent + caption-to-picture mapping runs too. Two for that, one for the blobs and + their text. Still flat in document count.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + for i in range(4): + document_id = f"doc-{i}" + await rag.document_item_repository.create_items( + document_id, + [ + DocumentItem( + document_id=document_id, + position=0, + self_ref="#/pictures/0", + label="picture", + text=f"caption for {document_id}", + picture_data=f"bytes-{document_id}".encode(), + ), + DocumentItem( + document_id=document_id, + position=1, + self_ref="#/texts/1", + label="caption", + text=f"figure 1 of {document_id}", + ), + ], + ) + + counts = [] + for n in (1, 4): + results = [_picture_result(f"doc-{i}", "#/texts/1") for i in range(n)] + item_queries["n"] = 0 + await _populate_image_data(rag, results) + counts.append(item_queries["n"]) + assert all(r.image_data for r in results) + + assert counts == [3, 3], counts From 5b0444043af807ed4fcc1657c9a4459f50202be9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 16:35:17 +0300 Subject: [PATCH 2/6] 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] From 460215158d0cde0c649370546af968c8c8a1c55a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 16:58:26 +0300 Subject: [PATCH 3/6] Batch the multimodal reranker's picture fetch `_attach_picture_data` fetched picture bytes once per document, over the `limit * 10` candidates reranking asks for, so it was the per-document fetch with the most candidates behind it. It now issues one query however many documents the candidates span: one for ten documents, as for one. Removes `get_text_for_refs`, whose only caller now gets the text back with the bytes from `get_pictures_grouped`. `test_client_search_include_images_false_skips_lookup` returned no search results, so asserting the picture accessor went uncalled held whatever the code did. It now returns a picture-carrying result, making "did not fetch" the assertion rather than "had nothing to fetch". --- CHANGELOG.md | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 14 ++++-- .../rag/store/repositories/document_item.py | 30 ----------- tests/store/test_document_items.py | 50 ------------------- tests/test_enrichment_batching.py | 45 ++++++++++++++++- tests/test_picture_in_context.py | 28 ++++++++--- 6 files changed, 74 insertions(+), 95 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ea0ee94..1dc425bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ### Changed -- 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`. +- Search enrichment, the multimodal reranker's picture fetch, 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`. `get_text_for_refs` is removed; `get_pictures_grouped` returns the text alongside the bytes. - 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 71bae916..8f8e9e0c 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -85,7 +85,11 @@ async def search( async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None: """Attach picture bytes to synthetic picture chunks in-place, so a multimodal reranker can score the pixels instead of just the chunk's - description text. Batches one picture-bytes lookup per document.""" + description text. + + One query however many documents the candidates span, which matters here + more than anywhere: reranking fetches `limit * 10` candidates. + """ by_doc: dict[str, list[tuple[Chunk, str]]] = {} for chunk in chunks: if chunk.document_id is None: @@ -94,11 +98,11 @@ async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None: if len(refs) == 1 and refs[0].startswith(PICTURE_REF_PREFIX): by_doc.setdefault(chunk.document_id, []).append((chunk, refs[0])) + bytes_by_document, _ = await client.document_item_repository.get_pictures_grouped( + {doc_id: [ref for _, ref in pairs] for doc_id, pairs in by_doc.items()} + ) for doc_id, doc_chunks in by_doc.items(): - refs = [ref for _, ref in doc_chunks] - bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk( - doc_id, refs - ) + bytes_by_ref = bytes_by_document.get(doc_id, {}) for chunk, ref in doc_chunks: data = bytes_by_ref.get(ref) if data: 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 7932868d..b8902cbf 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -394,36 +394,6 @@ class DocumentItemRepository: grouped.setdefault(row["document_id"], {})[caption] = row["self_ref"] return grouped - async def get_text_for_refs( - self, document_id: str, refs: list[str] - ) -> dict[str, str]: - """Fetch the ``text`` field for multiple self_refs within a single document. - - Returns ``{self_ref: text}`` for refs whose text is non-empty. Used - alongside ``get_pictures_for_chunk`` to label figures in agent-facing - search results: picture items carry their VLM-generated caption in - the ``text`` field, and the OpenAI vision message format has no - identifier on binary parts, so the caption text is the only signal a - model can use to correlate a description with the picture it sees. - """ - if not refs: - return {} - - safe_id = escape_sql_string(document_id) - refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) - rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "text"]) - .where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})") - .to_list() - ) - result: dict[str, str] = {} - for row in rows: - text = row.get("text") or "" - if text: - result[row["self_ref"]] = text - return result - async def get_caption_picture_refs( self, document_id: str, refs: list[str] ) -> dict[str, str]: diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index e6b1ed6c..cfdb0f69 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -607,56 +607,6 @@ class TestPictureDataStorage: # Empty refs returns empty dict assert await repo.get_pictures_for_chunk("doc-1", []) == {} - async def test_get_text_for_refs(self, temp_db_path): - """Text is returned for any ref with non-empty ``text``, regardless of label. - - In practice pictures carry their caption in the ``text`` field - (populated by the VLM picture-description pass during ingest); this - method surfaces that text alongside the picture bytes so the model can - correlate a description with the binary it sees. The same method also - returns text for non-picture refs — callers filter by label. - """ - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - await repo.create_items( - "doc-1", - [ - DocumentItem( - document_id="doc-1", - position=0, - self_ref="#/pictures/0", - label="picture", - text="Figure 1. CCS generation over time.", - picture_data=b"\x89PNG\r\n\x1a\nfake", - ), - DocumentItem( - document_id="doc-1", - position=1, - self_ref="#/pictures/1", - label="picture", - text="", # no VLM caption available - picture_data=b"\x89PNG\r\n\x1a\nfake2", - ), - DocumentItem( - document_id="doc-1", - position=2, - self_ref="#/texts/0", - label="paragraph", - text="Inline prose.", - ), - ], - ) - - captions = await repo.get_text_for_refs( - "doc-1", - ["#/pictures/0", "#/pictures/1", "#/texts/0", "#/pictures/999"], - ) - assert captions == { - "#/pictures/0": "Figure 1. CCS generation over time.", - "#/texts/0": "Inline prose.", - } - assert await repo.get_text_for_refs("doc-1", []) == {} - async def test_get_caption_picture_refs(self, temp_db_path): """A caption ref resolves to the picture at the immediately preceding position; a table caption (no preceding picture) resolves to nothing.""" diff --git a/tests/test_enrichment_batching.py b/tests/test_enrichment_batching.py index a5aa15c3..34af4a7d 100644 --- a/tests/test_enrichment_batching.py +++ b/tests/test_enrichment_batching.py @@ -2,8 +2,8 @@ import lancedb import pytest from haiku.rag.client import HaikuRAG -from haiku.rag.client.search import _populate_image_data -from haiku.rag.store.models import DocumentItem, SearchResult +from haiku.rag.client.search import _attach_picture_data, _populate_image_data +from haiku.rag.store.models import Chunk, DocumentItem, SearchResult def _picture_result(document_id: str, ref: str) -> SearchResult: @@ -210,3 +210,44 @@ async def test_expansion_widens_each_document_with_its_own_items(temp_db_path): 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"] + + +def _picture_chunk(document_id: str) -> Chunk: + return Chunk( + id=f"{document_id}-pic", + document_id=document_id, + content="a figure", + metadata={"doc_item_refs": ["#/pictures/0"], "labels": ["picture"]}, + ) + + +@pytest.mark.asyncio +async def test_reranker_blob_fetch_is_one_query_for_any_document_count( + temp_db_path, item_queries +): + """This path runs over `limit * 10` candidates, so per-document fetching + costs the most here.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + await _seed(rag, [f"doc-{i}" for i in range(10)]) + + counts = [] + for n in (1, 10): + chunks = [_picture_chunk(f"doc-{i}") for i in range(n)] + item_queries["n"] = 0 + await _attach_picture_data(rag, chunks) + counts.append(item_queries["n"]) + assert all(c._picture_data for c in chunks) + + assert counts == [1, 1], counts + + +@pytest.mark.asyncio +async def test_reranker_gives_each_chunk_its_own_document_picture(temp_db_path): + async with HaikuRAG(temp_db_path, create=True) as rag: + await _seed(rag, ["doc-a", "doc-b"]) + + chunks = [_picture_chunk("doc-a"), _picture_chunk("doc-b")] + await _attach_picture_data(rag, chunks) + + assert chunks[0]._picture_data == b"bytes-doc-a" + assert chunks[1]._picture_data == b"bytes-doc-b" diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 944afd56..c2c4c497 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -16,7 +16,7 @@ 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.store.models.chunk import SearchResult +from haiku.rag.store.models.chunk import Chunk, 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 @@ -138,20 +138,34 @@ async def test_client_search_include_images_false_skips_lookup(temp_db_path): ], ) # Spy that we never reach the picture-bytes accessor - rag.document_item_repository.get_pictures_for_chunk = AsyncMock( # type: ignore[method-assign] - wraps=rag.document_item_repository.get_pictures_for_chunk + rag.document_item_repository.get_pictures_grouped = AsyncMock( # type: ignore[method-assign] + wraps=rag.document_item_repository.get_pictures_grouped ) from haiku.rag.client.search import search - # Stub the chunk-search results so we don't depend on embeddings/FTS + # A real picture-carrying result, so not fetching is the assertion + # rather than there being nothing to fetch. async def fake_chunk_search(*args, **kwargs): - return [] + return [ + ( + Chunk( + id="chunk-1", + document_id="doc-1", + content="body", + metadata={"doc_item_refs": ["#/pictures/0"]}, + ), + 0.9, + ) + ] rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign] - await search(rag, "anything", include_images=False) - rag.document_item_repository.get_pictures_for_chunk.assert_not_called() + results = await search(rag, "anything", include_images=False) + + assert len(results) == 1 + assert results[0].image_data is None + rag.document_item_repository.get_pictures_grouped.assert_not_called() @pytest.mark.asyncio From 65f6d72cd5d373af91ca3c78cae96fb48d3b9ec6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 17:12:10 +0300 Subject: [PATCH 4/6] Remove the single-document item accessors the batching replaced `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs` and `get_all_items_grouped` have no callers left: the grouped equivalents serve every path that used them. `get_all_items_grouped` had none even before this branch. Tests whose subject was a removed method go with it. Tests that only used one to fetch a fixture now use the grouped call, so what they assert is unchanged. --- CHANGELOG.md | 2 +- .../rag/store/repositories/document_item.py | 109 +---------------- tests/store/test_document_items.py | 115 ++++-------------- tests/test_context.py | 4 +- tests/test_context_enhancement.py | 4 +- 5 files changed, 30 insertions(+), 204 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1dc425bd..e84e192d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ### Changed -- Search enrichment, the multimodal reranker's picture fetch, 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`. `get_text_for_refs` is removed; `get_pictures_grouped` returns the text alongside the bytes. +- Search enrichment, the multimodal reranker's picture fetch, 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`. The single-document methods they replace are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns the text alongside the bytes. - 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/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index b8902cbf..0426801b 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -108,69 +108,10 @@ class DocumentItemRepository: items.sort(key=lambda x: x.position) return items - async def get_all_items_grouped( - self, document_ids: list[str] | None = None - ) -> dict[str, list[DocumentItem]]: - """Get all items grouped by document_id in a single query. - - Args: - document_ids: If provided, only fetch items for these documents. - If None, fetches all items. - - Returns: - Dict mapping document_id to sorted list of DocumentItem. - """ - query = self.store.document_items_table.query().select(_METADATA_COLUMNS) - if document_ids is not None: - safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids) - query = query.where(f"document_id IN ({safe_ids})") - rows = await query.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_items_in_range( - self, document_id: str, start: int, end: int - ) -> list[DocumentItem]: - """Get items for a document within a position range (inclusive).""" - safe_id = escape_sql_string(document_id) - rows = await ( - self.store.document_items_table.query() - .select(_METADATA_COLUMNS) - .where( - f"document_id = '{safe_id}' " - f"AND position >= {start} AND position <= {end}" - ) - .to_list() - ) - items = [self._record_to_item(row) for row in rows] - items.sort(key=lambda x: x.position) - return items - - async def resolve_refs(self, document_id: str, refs: list[str]) -> dict[str, int]: - """Resolve self_refs to positions. Returns {self_ref: position}.""" - if not refs: - return {} - - safe_id = escape_sql_string(document_id) - refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) - rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "position"]) - .where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})") - .to_list() - ) - 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.""" + """Resolve self_refs to positions, across documents, in one query.""" predicate = self._per_document_predicate(refs_by_document, "self_ref") if predicate is None: return {} @@ -190,7 +131,7 @@ class DocumentItemRepository: 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. + """Items within a position range per document, 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 @@ -352,7 +293,7 @@ class DocumentItemRepository: async def get_caption_picture_refs_grouped( self, refs_by_document: "Mapping[str, list[str]]" ) -> dict[str, dict[str, str]]: - """`get_caption_picture_refs` across documents in two queries. + """Map caption refs to the picture preceding them, in two queries. Two rather than one because the stages are dependent: a caption's picture is the item at `position - 1`, which the first query is what @@ -393,47 +334,3 @@ class DocumentItemRepository: if caption: grouped.setdefault(row["document_id"], {})[caption] = row["self_ref"] return grouped - - async def get_caption_picture_refs( - self, document_id: str, refs: list[str] - ) -> dict[str, str]: - """Map caption refs to the picture item immediately preceding them. - - Docling emits a figure's caption at the position right after its - picture, so a caption's picture is the picture item at - ``position - 1``. Returns ``{caption_ref: picture_ref}`` for the - caption refs among ``refs`` that have a picture predecessor. Non-caption - refs, and captions whose predecessor is not a picture (table captions), - map to nothing. - """ - if not refs: - return {} - - safe_id = escape_sql_string(document_id) - refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs) - caption_rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "position"]) - .where( - f"document_id = '{safe_id}' AND label = 'caption' " - f"AND self_ref IN ({refs_sql})" - ) - .to_list() - ) - if not caption_rows: - return {} - - prev_to_caption = {row["position"] - 1: row["self_ref"] for row in caption_rows} - positions_sql = ", ".join(str(p) for p in prev_to_caption) - picture_rows = await ( - self.store.document_items_table.query() - .select(["self_ref", "position"]) - .where( - f"document_id = '{safe_id}' AND label = 'picture' " - f"AND position IN ({positions_sql})" - ) - .to_list() - ) - return { - prev_to_caption[row["position"]]: row["self_ref"] for row in picture_rows - } diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index cfdb0f69..6104f2fc 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -265,33 +265,14 @@ class TestDocumentItemRepository: ] await repo.create_items("doc-1", items) - result = await repo.get_items_in_range("doc-1", 3, 7) + result = (await repo.get_items_in_ranges({"doc-1": (3, 7)})).get( + "doc-1", [] + ) assert len(result) == 5 assert result[0].position == 3 assert result[-1].position == 7 assert result[0].text == "Item 3" - async def test_resolve_refs(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - - items = [ - DocumentItem( - document_id="doc-1", - position=i, - self_ref=f"#/texts/{i}", - label="paragraph", - text=f"Item {i}", - ) - for i in range(10) - ] - await repo.create_items("doc-1", items) - - refs = await repo.resolve_refs( - "doc-1", ["#/texts/2", "#/texts/7", "#/texts/999"] - ) - assert refs == {"#/texts/2": 2, "#/texts/7": 7} - async def test_get_item_count(self, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as rag: repo = DocumentItemRepository(rag.store) @@ -376,25 +357,24 @@ class TestDocumentItemRepository: (0, 2), ] - in_range = await repo.get_items_in_range("doc-1", 0, 2) + in_range = (await repo.get_items_in_ranges({"doc-1": (0, 2)})).get( + "doc-1", [] + ) assert [(i.heading_level, i.tree_depth) for i in in_range] == [ (1, 1), (2, 2), (0, 2), ] - grouped = await repo.get_all_items_grouped(["doc-1"]) - assert [(i.heading_level, i.tree_depth) for i in grouped["doc-1"]] == [ + assert [ + (i.heading_level, i.tree_depth) + for i in await repo.get_all_items("doc-1") + ] == [ (1, 1), (2, 2), (0, 2), ] - async def test_empty_refs_returns_empty(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - assert await repo.resolve_refs("doc-1", []) == {} - async def test_items_sorted_by_position(self, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as rag: repo = DocumentItemRepository(rag.store) @@ -412,7 +392,9 @@ class TestDocumentItemRepository: ] await repo.create_items("doc-1", items) - result = await repo.get_items_in_range("doc-1", 0, 9) + result = (await repo.get_items_in_ranges({"doc-1": (0, 9)})).get( + "doc-1", [] + ) positions = [item.position for item in result] assert positions == sorted(positions) @@ -440,9 +422,11 @@ class TestDocumentItemPopulation: count = await rag.document_item_repository.get_item_count(created.id) assert count == 6 - items = await rag.document_item_repository.get_items_in_range( - created.id, 0, count - ) + items = ( + await rag.document_item_repository.get_items_in_ranges( + {created.id: (0, count)} + ) + ).get(created.id, []) assert items[0].label == "section_header" assert items[0].text == "Introduction" assert items[1].label == "paragraph" @@ -607,62 +591,6 @@ class TestPictureDataStorage: # Empty refs returns empty dict assert await repo.get_pictures_for_chunk("doc-1", []) == {} - async def test_get_caption_picture_refs(self, temp_db_path): - """A caption ref resolves to the picture at the immediately preceding - position; a table caption (no preceding picture) resolves to nothing.""" - async with HaikuRAG(temp_db_path, create=True) as rag: - repo = DocumentItemRepository(rag.store) - await repo.create_items( - "doc-1", - [ - DocumentItem( - document_id="doc-1", - position=0, - self_ref="#/pictures/0", - label="picture", - picture_data=b"\x89PNG\r\n\x1a\nfake", - ), - DocumentItem( - document_id="doc-1", - position=1, - self_ref="#/texts/0", - label="caption", - text="Figure 1. A figure caption.", - ), - DocumentItem( - document_id="doc-1", - position=2, - self_ref="#/texts/1", - label="paragraph", - text="Body prose.", - ), - DocumentItem( - document_id="doc-1", - position=3, - self_ref="#/tables/0", - label="table", - text="| a | b |", - ), - DocumentItem( - document_id="doc-1", - position=4, - self_ref="#/texts/2", - label="caption", - text="Table 1. A table caption.", - ), - ], - ) - - # Figure caption resolves to its picture; table caption does not. - got = await repo.get_caption_picture_refs( - "doc-1", ["#/texts/0", "#/texts/1", "#/texts/2"] - ) - assert got == {"#/texts/0": "#/pictures/0"} - - # A non-caption ref alone yields nothing. - assert await repo.get_caption_picture_refs("doc-1", ["#/texts/1"]) == {} - assert await repo.get_caption_picture_refs("doc-1", []) == {} - async def test_hot_paths_exclude_picture_data(self, temp_db_path): """Light read paths must NOT pull picture_data into memory.""" async with HaikuRAG(temp_db_path, create=True) as rag: @@ -685,10 +613,9 @@ class TestPictureDataStorage: for item in await repo.get_all_items("doc-1"): assert item.picture_data is None - for item in await repo.get_items_in_range("doc-1", 0, 10): - assert item.picture_data is None - grouped = await repo.get_all_items_grouped(["doc-1"]) - for item in grouped["doc-1"]: + for item in (await repo.get_items_in_ranges({"doc-1": (0, 10)})).get( + "doc-1", [] + ): assert item.picture_data is None # But the picture-byte accessors still work diff --git a/tests/test_context.py b/tests/test_context.py index 2d73c792..48130fdd 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -1462,10 +1462,10 @@ class TestExpandWithItemsWindowEdges: ) async def no_window(*_args, **_kwargs): - return [] + return {} monkeypatch.setattr( - rag.document_item_repository, "get_items_in_range", no_window + rag.document_item_repository, "get_items_in_ranges", no_window ) result = SearchResult( diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index 6696edf7..9fa11e72 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -322,7 +322,9 @@ async def test_expand_context_single_item_document(temp_db_path): assert doc.id is not None # Create a search result with a doc_item_ref pointing to the item - items = await client.document_item_repository.get_items_in_range(doc.id, 0, 10) + items = ( + await client.document_item_repository.get_items_in_ranges({doc.id: (0, 10)}) + ).get(doc.id, []) assert len(items) > 0 search_results = [ From 28217fcf824966dcfa6d1d9e87df3716da8c7d0a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 17:41:10 +0300 Subject: [PATCH 5/6] Keep result order when expansion is batched Splitting the assembly into a passthrough pass and an expandable pass reordered equal-scored results: the score sort that follows is stable, so the order results arrive in is the tiebreak. Results are assembled in document_groups order again, after the batched fetch rather than around it. Also ports the caption negative cases the removed single-document test carried: a table's caption and an ordinary text reference map to no picture. --- CHANGELOG.md | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 12 ++++--- tests/store/test_document_item_grouped.py | 42 +++++++++++++++++++++++ tests/test_enrichment_batching.py | 22 ++++++++++++ 4 files changed, 72 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e84e192d..cdd0613c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ### Changed -- Search enrichment, the multimodal reranker's picture fetch, 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`. The single-document methods they replace are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns the text alongside the bytes. +- Search enrichment, the multimodal reranker's picture fetch, 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`. The superseded methods are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns the text alongside the bytes. - 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 8f8e9e0c..4bd1dc0d 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -237,10 +237,6 @@ async def expand_context( 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 not in expandable: - expanded_results.extend(doc_results) - repo = client.document_item_repository positions_by_document = await repo.resolve_refs_grouped( { @@ -255,7 +251,13 @@ async def expand_context( } items_by_document = await repo.get_items_in_ranges(windows) - for doc_id, doc_results in expandable.items(): + # In document_groups order: the score sort below is stable, so assembling + # expandable and passthrough documents in separate passes would reorder + # equal-scored results. + for doc_id, doc_results in document_groups.items(): + if doc_id not in expandable: + expanded_results.extend(doc_results) + continue expanded_results.extend( expand_with_items( doc_results, diff --git a/tests/store/test_document_item_grouped.py b/tests/store/test_document_item_grouped.py index 2daf03c7..9a975bab 100644 --- a/tests/store/test_document_item_grouped.py +++ b/tests/store/test_document_item_grouped.py @@ -114,3 +114,45 @@ async def test_grouped_calls_with_nothing_asked_for_do_not_query( assert await repo.get_caption_picture_refs_grouped({}) == {} assert item_queries["n"] == 0 + + +@pytest.mark.asyncio +async def test_caption_picture_refs_grouped_ignores_non_picture_predecessors( + temp_db_path, +): + """A caption maps to a picture only. A table's caption, and an ordinary text + reference, map to nothing.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + await repo.create_items( + "doc-1", + [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/tables/0", + label="table", + text="a table", + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/table-caption", + label="caption", + text="Table 1", + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/texts/plain", + label="text", + text="ordinary prose", + ), + ], + ) + + got = await repo.get_caption_picture_refs_grouped( + {"doc-1": ["#/texts/table-caption", "#/texts/plain"]} + ) + + assert got == {} diff --git a/tests/test_enrichment_batching.py b/tests/test_enrichment_batching.py index 34af4a7d..3c13abd4 100644 --- a/tests/test_enrichment_batching.py +++ b/tests/test_enrichment_batching.py @@ -251,3 +251,25 @@ async def test_reranker_gives_each_chunk_its_own_document_picture(temp_db_path): assert chunks[0]._picture_data == b"bytes-doc-a" assert chunks[1]._picture_data == b"bytes-doc-b" + + +@pytest.mark.asyncio +async def test_expansion_keeps_document_order_for_tied_scores(temp_db_path): + """The score sort is stable, so equal-scored results must come back in the + order they arrived, whether or not their document expands.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + await _seed_expandable(rag, ["doc-expandable"]) + + passthrough = SearchResult( + chunk_id="doc-plain-anchor", + document_id="doc-plain", + content="plain body", + score=0.5, + doc_item_refs=[], + ) + expandable = _text_result("doc-expandable") + expandable.score = 0.5 + + for order in ([passthrough, expandable], [expandable, passthrough]): + expanded = await rag.expand_context(list(order)) + assert [r.chunk_id for r in expanded] == [r.chunk_id for r in order] From 62da6086b8d9a804ebf3d42854973d9d67145c92 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 18:00:00 +0300 Subject: [PATCH 6/6] Stop the reranker fetch reading text it discards Collapsing the caption text into `get_pictures_grouped` served the enrichment path, which uses it, but the multimodal reranker discards the second return value while still paying to read the column. That is the widest fan-out in the codebase, `limit * 10` candidates, and it previously projected self_ref and picture_data alone. `with_text` is opt-in and off by default, so the cheap projection is what a caller gets unless it asks for more. The reranker test asserts the projection as well as the query count, since a count alone would not notice the column coming back. --- CHANGELOG.md | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 4 +++- .../rag/store/repositories/document_item.py | 17 +++++++++----- tests/store/test_document_item_grouped.py | 17 +++++++++++++- tests/test_enrichment_batching.py | 22 ++++++++++++++++++- 5 files changed, 53 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cdd0613c..1daa4601 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ ### Changed -- Search enrichment, the multimodal reranker's picture fetch, 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`. The superseded methods are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns the text alongside the bytes. +- Search enrichment, the multimodal reranker's picture fetch, 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`. The superseded methods are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns each picture's text alongside its bytes under `with_text`, off by default so the reranker's blob fetch does not read a column it discards. - 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 4bd1dc0d..a85e1737 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -184,7 +184,9 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) if not wanted: return - bytes_by_document, captions_by_document = await repo.get_pictures_grouped(wanted) + bytes_by_document, captions_by_document = await repo.get_pictures_grouped( + wanted, with_text=True + ) if not bytes_by_document: return 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 0426801b..4f293498 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -260,21 +260,28 @@ class DocumentItemRepository: return " OR ".join(clauses) if clauses else None async def get_pictures_grouped( - self, refs_by_document: "Mapping[str, list[str]]" + self, + refs_by_document: "Mapping[str, list[str]]", + *, + with_text: bool = False, ) -> tuple[dict[str, dict[str, bytes]], dict[str, dict[str, str]]]: - """Picture bytes and their text, across documents, in one query. + """Picture bytes across documents in one query. Returns `(bytes_by_document, text_by_document)`, each `{document_id: {self_ref: value}}` and each omitting refs whose value is - empty. The text comes from the same rows as the bytes, so asking for it - separately would be a second read of rows already in hand. + empty. `with_text` adds each picture's text to the projection, which is + free in queries because it is on the same rows, but not in bytes: the + text column is dead weight for a caller that only scores pixels. """ predicate = self._per_document_predicate(refs_by_document, "self_ref") if predicate is None: return {}, {} + columns = ["document_id", "self_ref", "picture_data"] + if with_text: + columns.append("text") rows = await ( self.store.document_items_table.query() - .select(["document_id", "self_ref", "picture_data", "text"]) + .select(columns) .where(predicate) .to_list() ) diff --git a/tests/store/test_document_item_grouped.py b/tests/store/test_document_item_grouped.py index 9a975bab..ae8d04f9 100644 --- a/tests/store/test_document_item_grouped.py +++ b/tests/store/test_document_item_grouped.py @@ -54,7 +54,7 @@ async def test_pictures_grouped_keeps_documents_apart(temp_db_path, item_queries item_queries["n"] = 0 blobs, texts = await repo.get_pictures_grouped( - {"doc-a": ["#/pictures/0"], "doc-b": ["#/pictures/0"]} + {"doc-a": ["#/pictures/0"], "doc-b": ["#/pictures/0"]}, with_text=True ) assert item_queries["n"] == 1 @@ -156,3 +156,18 @@ async def test_caption_picture_refs_grouped_ignores_non_picture_predecessors( ) assert got == {} + + +@pytest.mark.asyncio +async def test_pictures_grouped_omits_text_unless_asked(temp_db_path, item_queries): + """Text is dead weight for a caller that only scores pixels.""" + async with Store(temp_db_path, create=True) as store: + repo = DocumentItemRepository(store) + await _seed(repo, "doc-a") + + item_queries["n"] = 0 + blobs, texts = await repo.get_pictures_grouped({"doc-a": ["#/pictures/0"]}) + + assert item_queries["n"] == 1 + assert blobs == {"doc-a": {"#/pictures/0": b"bytes-doc-a"}} + assert texts == {} diff --git a/tests/test_enrichment_batching.py b/tests/test_enrichment_batching.py index 3c13abd4..0933e720 100644 --- a/tests/test_enrichment_batching.py +++ b/tests/test_enrichment_batching.py @@ -1,4 +1,5 @@ import lancedb +import lancedb.query import pytest from haiku.rag.client import HaikuRAG @@ -35,6 +36,21 @@ async def _seed(rag: HaikuRAG, document_ids: list[str]) -> None: ) +@pytest.fixture +def item_projections(monkeypatch): + """Columns each document_items query projects.""" + projections: list[list[str]] = [] + select = lancedb.query.AsyncQuery.select + + def recording(self, columns): + if isinstance(columns, list): + projections.append([str(c) for c in columns]) + return select(self, columns) + + monkeypatch.setattr(lancedb.query.AsyncQuery, "select", recording) + return projections + + @pytest.fixture def item_queries(monkeypatch): tally = {"n": 0} @@ -223,7 +239,7 @@ def _picture_chunk(document_id: str) -> Chunk: @pytest.mark.asyncio async def test_reranker_blob_fetch_is_one_query_for_any_document_count( - temp_db_path, item_queries + temp_db_path, item_queries, item_projections ): """This path runs over `limit * 10` candidates, so per-document fetching costs the most here.""" @@ -239,6 +255,10 @@ async def test_reranker_blob_fetch_is_one_query_for_any_document_count( assert all(c._picture_data for c in chunks) assert counts == [1, 1], counts + # The reranker scores pixels, so `text` has no business in the projection. + picture_projections = [p for p in item_projections if "picture_data" in p] + assert picture_projections, "no picture query observed" + assert all("text" not in p for p in picture_projections), picture_projections @pytest.mark.asyncio