danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Batch search enrichment across documents

Browse source 
`_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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-18 15:50:42 +03:00
parent d177b5883b
commit af6a6b0bbe
No known key found for this signature in database
5 changed files with 388 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

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

68
haiku_rag_slim/haiku/rag/client/search.py
View file

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

100
haiku_rag_slim/haiku/rag/store/repositories/document_item.py
View file

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

116
tests/store/test_document_item_grouped.py Normal file
View file

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

134
tests/test_enrichment_batching.py Normal file
View file

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