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

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

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 = bytes_by_document.get(doc_id, {})
-        bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk(
-            doc_id, refs
-        )
         for chunk, ref in doc_chunks:
             data = bytes_by_ref.get(ref)
             if data:

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]:

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

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.client.search import _attach_picture_data, _populate_image_data
-from haiku.rag.store.models import DocumentItem, SearchResult
+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"

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]
+        rag.document_item_repository.get_pictures_grouped = AsyncMock(  # type: ignore[method-assign]
-            wraps=rag.document_item_repository.get_pictures_for_chunk
+            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)
+        results = await search(rag, "anything", include_images=False)
-        rag.document_item_repository.get_pictures_for_chunk.assert_not_called()
+
+        assert len(results) == 1
+        assert results[0].image_data is None
+        rag.document_item_repository.get_pictures_grouped.assert_not_called()
 
 
 @pytest.mark.asyncio