From 62da6086b8d9a804ebf3d42854973d9d67145c92 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Aug 2026 18:00:00 +0300 Subject: [PATCH] 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