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