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