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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-18 18:00:00 +03:00
parent 28217fcf82
commit 62da6086b8
No known key found for this signature in database
5 changed files with 53 additions and 9 deletions

View file

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

View file

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

View file

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

View file

@ -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 == {}

View file

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