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

Merge pull request #555 from ggozad/feat/batch-enrichment

Browse source 
Batch document-item fetches across documents
This commit is contained in:
Yiorgis Gozadinos 2026-08-19 11:04:15 +03:00 committed by GitHub
commit 13bd69b908
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 780 additions and 356 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, 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.

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

@ -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:
@ -139,22 +143,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 +179,33 @@ 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, with_text=True
)
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(
@ -209,7 +221,7 @@ async def expand_context(
chunks were created without docling metadata (e.g., custom chunks passed
to import_document).
"""
from haiku.rag.context import expand_with_items
from haiku.rag.context import expand_with_items, window_for
max_chars = client._config.search.max_context_chars
@ -222,24 +234,40 @@ async def expand_context(
document_groups[doc_id].append(result)
expanded_results = []
expandable = {
doc_id: doc_results
for doc_id, doc_results in document_groups.items()
if doc_id is not None and any(r.doc_item_refs for r in doc_results)
}
repo = client.document_item_repository
positions_by_document = await repo.resolve_refs_grouped(
{
doc_id: [ref for r in doc_results for ref in r.doc_item_refs]
for doc_id, doc_results in expandable.items()
}
)
windows = {
doc_id: window_for(positions)
for doc_id, positions in positions_by_document.items()
if positions
}
items_by_document = await repo.get_items_in_ranges(windows)
# In document_groups order: the score sort below is stable, so assembling
# expandable and passthrough documents in separate passes would reorder
# equal-scored results.
for doc_id, doc_results in document_groups.items():
if doc_id is None:
if doc_id not in expandable:
expanded_results.extend(doc_results)
continue
has_refs = any(r.doc_item_refs for r in doc_results)
if not has_refs:
expanded_results.extend(doc_results)
continue
expanded = await expand_with_items(
client.document_item_repository,
doc_id,
doc_results,
max_chars,
expanded_results.extend(
expand_with_items(
doc_results,
max_chars,
positions_by_document.get(doc_id, {}),
items_by_document.get(doc_id, []),
)
)
expanded_results.extend(expanded)
expanded_results.sort(key=lambda r: r.score, reverse=True)
# image_data and picture_captions are preserved through expansion by

42
haiku_rag_slim/haiku/rag/context.py
View file

@ -34,7 +34,6 @@ In both cases:
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.store.repositories.document_item import DocumentItemRepository
_NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
_SECTION_BOUNDARY_LABELS = {"section_header", "title"}
@ -425,33 +424,28 @@ def _build_result(
_WINDOW_MARGIN = 100
async def expand_with_items(
document_item_repository: DocumentItemRepository,
document_id: str,
def window_for(ref_positions: dict[str, int]) -> tuple[int, int]:
"""The inclusive position range to fetch around a document's matches.
The margin must be wide enough to find section boundaries: the nearest
section_header or title above and below the match.
"""
positions = sorted(ref_positions.values())
return max(0, positions[0] - _WINDOW_MARGIN), positions[-1] + _WINDOW_MARGIN
def expand_with_items(
results: list[SearchResult],
max_chars: int,
ref_positions: dict[str, int],
window_items: list[DocumentItem],
) -> list[SearchResult]:
"""Expand results using the document_items table."""
all_refs = []
for result in results:
all_refs.extend(result.doc_item_refs)
"""Expand results from items already fetched.
ref_positions = await document_item_repository.resolve_refs(document_id, all_refs)
if not ref_positions:
return results
# Fetch a window of items around matched positions. The margin must be
# wide enough to find section boundaries (the nearest section_header/title
# above and below the match).
all_positions = sorted(ref_positions.values())
window_margin = _WINDOW_MARGIN
window_start = max(0, min(all_positions) - window_margin)
window_end = max(all_positions) + window_margin
window_items = await document_item_repository.get_items_in_range(
document_id, window_start, window_end
)
if not window_items:
Fetching is the caller's, so one query can serve every document in a result
set rather than one per document.
"""
if not ref_positions or not window_items:
return results
has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items)

220
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
@ -107,24 +108,50 @@ class DocumentItemRepository:
items.sort(key=lambda x: x.position)
return items
async def get_all_items_grouped(
self, document_ids: list[str] | None = None
async def resolve_refs_grouped(
self, refs_by_document: "Mapping[str, Sequence[str]]"
) -> dict[str, dict[str, int]]:
"""Resolve self_refs to positions, across documents, in one query."""
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", "position"])
.where(predicate)
.to_list()
)
grouped: dict[str, dict[str, int]] = {}
for row in rows:
grouped.setdefault(row["document_id"], {})[row["self_ref"]] = row[
"position"
]
return grouped
async def get_items_in_ranges(
self, ranges_by_document: "Mapping[str, tuple[int, int]]"
) -> dict[str, list[DocumentItem]]:
"""Get all items grouped by document_id in a single query.
"""Items within a position range per document, in one query.
Args:
document_ids: If provided, only fetch items for these documents.
If None, fetches all items.
Returns:
Dict mapping document_id to sorted list of DocumentItem.
Each document keeps its own inclusive range. Positions repeat across
documents, so a shared range would splice one document's items into
another's context.
"""
query = self.store.document_items_table.query().select(_METADATA_COLUMNS)
if document_ids is not None:
safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
query = query.where(f"document_id IN ({safe_ids})")
rows = await query.to_list()
clauses = []
for document_id, (start, end) in ranges_by_document.items():
safe_id = escape_sql_string(document_id)
clauses.append(
f"(document_id = '{safe_id}' "
f"AND position >= {start} AND position <= {end})"
)
if not clauses:
return {}
rows = await (
self.store.document_items_table.query()
.select(_METADATA_COLUMNS)
.where(" OR ".join(clauses))
.to_list()
)
grouped: dict[str, list[DocumentItem]] = {}
for row in rows:
item = self._record_to_item(row)
@ -133,39 +160,6 @@ class DocumentItemRepository:
items.sort(key=lambda x: x.position)
return grouped
async def get_items_in_range(
self, document_id: str, start: int, end: int
) -> list[DocumentItem]:
"""Get items for a document within a position range (inclusive)."""
safe_id = escape_sql_string(document_id)
rows = await (
self.store.document_items_table.query()
.select(_METADATA_COLUMNS)
.where(
f"document_id = '{safe_id}' "
f"AND position >= {start} AND position <= {end}"
)
.to_list()
)
items = [self._record_to_item(row) for row in rows]
items.sort(key=lambda x: x.position)
return items
async def resolve_refs(self, document_id: str, refs: list[str]) -> dict[str, int]:
"""Resolve self_refs to positions. Returns {self_ref: position}."""
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", "position"])
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
.to_list()
)
return {row["self_ref"]: row["position"] for row in rows}
async def get_item_count(self, document_id: str) -> int:
"""Count items for a document."""
safe_id = escape_sql_string(document_id)
@ -242,76 +236,108 @@ class DocumentItemRepository:
result[row["self_ref"]] = data
return result
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.
@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 ...)`.
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.
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.
"""
if not refs:
return {}
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
safe_id = escape_sql_string(document_id)
refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
async def get_pictures_grouped(
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 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. `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(["self_ref", "text"])
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
.select(columns)
.where(predicate)
.to_list()
)
result: dict[str, str] = {}
blobs: dict[str, dict[str, bytes]] = {}
texts: dict[str, dict[str, str]] = {}
for row in rows:
text = row.get("text") or ""
data = row.get("picture_data")
if not data:
continue
blobs.setdefault(row["document_id"], {})[row["self_ref"]] = data
text = row.get("text")
if text:
result[row["self_ref"]] = text
return result
texts.setdefault(row["document_id"], {})[row["self_ref"]] = text
return blobs, texts
async def get_caption_picture_refs(
self, document_id: str, refs: list[str]
) -> dict[str, str]:
"""Map caption refs to the picture item immediately preceding them.
async def get_caption_picture_refs_grouped(
self, refs_by_document: "Mapping[str, list[str]]"
) -> dict[str, dict[str, str]]:
"""Map caption refs to the picture preceding them, in two queries.
Docling emits a figure's caption at the position right after its
picture, so a caption's picture is the picture item at
``position - 1``. Returns ``{caption_ref: picture_ref}`` for the
caption refs among ``refs`` that have a picture predecessor. Non-caption
refs, and captions whose predecessor is not a picture (table captions),
map to nothing.
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.
"""
if not refs:
predicate = self._per_document_predicate(refs_by_document, "self_ref")
if predicate is None:
return {}
safe_id = escape_sql_string(document_id)
refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
caption_rows = await (
self.store.document_items_table.query()
.select(["self_ref", "position"])
.where(
f"document_id = '{safe_id}' AND label = 'caption' "
f"AND self_ref IN ({refs_sql})"
)
.select(["document_id", "self_ref", "position"])
.where(f"label = 'caption' AND ({predicate})")
.to_list()
)
if not caption_rows:
return {}
prev_to_caption = {row["position"] - 1: row["self_ref"] for row in caption_rows}
positions_sql = ", ".join(str(p) for p in prev_to_caption)
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(["self_ref", "position"])
.where(
f"document_id = '{safe_id}' AND label = 'picture' "
f"AND position IN ({positions_sql})"
)
.select(["document_id", "self_ref", "position"])
.where(f"label = 'picture' AND ({picture_predicate})")
.to_list()
)
return {
prev_to_caption[row["position"]]: row["self_ref"] for row in picture_rows
}
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

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

@ -0,0 +1,173 @@
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"]}, with_text=True
)
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
@pytest.mark.asyncio
async def test_caption_picture_refs_grouped_ignores_non_picture_predecessors(
temp_db_path,
):
"""A caption maps to a picture only. A table's caption, and an ordinary text
reference, map to nothing."""
async with Store(temp_db_path, create=True) as store:
repo = DocumentItemRepository(store)
await repo.create_items(
"doc-1",
[
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/tables/0",
label="table",
text="a table",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/table-caption",
label="caption",
text="Table 1",
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/plain",
label="text",
text="ordinary prose",
),
],
)
got = await repo.get_caption_picture_refs_grouped(
{"doc-1": ["#/texts/table-caption", "#/texts/plain"]}
)
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 == {}

165
tests/store/test_document_items.py
View file

@ -265,33 +265,14 @@ class TestDocumentItemRepository:
]
await repo.create_items("doc-1", items)
result = await repo.get_items_in_range("doc-1", 3, 7)
result = (await repo.get_items_in_ranges({"doc-1": (3, 7)})).get(
"doc-1", []
)
assert len(result) == 5
assert result[0].position == 3
assert result[-1].position == 7
assert result[0].text == "Item 3"
async def test_resolve_refs(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(10)
]
await repo.create_items("doc-1", items)
refs = await repo.resolve_refs(
"doc-1", ["#/texts/2", "#/texts/7", "#/texts/999"]
)
assert refs == {"#/texts/2": 2, "#/texts/7": 7}
async def test_get_item_count(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
@ -376,25 +357,24 @@ class TestDocumentItemRepository:
(0, 2),
]
in_range = await repo.get_items_in_range("doc-1", 0, 2)
in_range = (await repo.get_items_in_ranges({"doc-1": (0, 2)})).get(
"doc-1", []
)
assert [(i.heading_level, i.tree_depth) for i in in_range] == [
(1, 1),
(2, 2),
(0, 2),
]
grouped = await repo.get_all_items_grouped(["doc-1"])
assert [(i.heading_level, i.tree_depth) for i in grouped["doc-1"]] == [
assert [
(i.heading_level, i.tree_depth)
for i in await repo.get_all_items("doc-1")
] == [
(1, 1),
(2, 2),
(0, 2),
]
async def test_empty_refs_returns_empty(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
assert await repo.resolve_refs("doc-1", []) == {}
async def test_items_sorted_by_position(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
@ -412,7 +392,9 @@ class TestDocumentItemRepository:
]
await repo.create_items("doc-1", items)
result = await repo.get_items_in_range("doc-1", 0, 9)
result = (await repo.get_items_in_ranges({"doc-1": (0, 9)})).get(
"doc-1", []
)
positions = [item.position for item in result]
assert positions == sorted(positions)
@ -440,9 +422,11 @@ class TestDocumentItemPopulation:
count = await rag.document_item_repository.get_item_count(created.id)
assert count == 6
items = await rag.document_item_repository.get_items_in_range(
created.id, 0, count
)
items = (
await rag.document_item_repository.get_items_in_ranges(
{created.id: (0, count)}
)
).get(created.id, [])
assert items[0].label == "section_header"
assert items[0].text == "Introduction"
assert items[1].label == "paragraph"
@ -607,112 +591,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."""
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",
picture_data=b"\x89PNG\r\n\x1a\nfake",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/0",
label="caption",
text="Figure 1. A figure caption.",
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/1",
label="paragraph",
text="Body prose.",
),
DocumentItem(
document_id="doc-1",
position=3,
self_ref="#/tables/0",
label="table",
text="| a | b |",
),
DocumentItem(
document_id="doc-1",
position=4,
self_ref="#/texts/2",
label="caption",
text="Table 1. A table caption.",
),
],
)
# Figure caption resolves to its picture; table caption does not.
got = await repo.get_caption_picture_refs(
"doc-1", ["#/texts/0", "#/texts/1", "#/texts/2"]
)
assert got == {"#/texts/0": "#/pictures/0"}
# A non-caption ref alone yields nothing.
assert await repo.get_caption_picture_refs("doc-1", ["#/texts/1"]) == {}
assert await repo.get_caption_picture_refs("doc-1", []) == {}
async def test_hot_paths_exclude_picture_data(self, temp_db_path):
"""Light read paths must NOT pull picture_data into memory."""
async with HaikuRAG(temp_db_path, create=True) as rag:
@ -735,10 +613,9 @@ class TestPictureDataStorage:
for item in await repo.get_all_items("doc-1"):
assert item.picture_data is None
for item in await repo.get_items_in_range("doc-1", 0, 10):
assert item.picture_data is None
grouped = await repo.get_all_items_grouped(["doc-1"])
for item in grouped["doc-1"]:
for item in (await repo.get_items_in_ranges({"doc-1": (0, 10)})).get(
"doc-1", []
):
assert item.picture_data is None
# But the picture-byte accessors still work

71
tests/test_context.py
View file

@ -8,6 +8,7 @@ from haiku.rag.context import (
_find_expansion_range,
_merge_ranges,
expand_with_items,
window_for,
)
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
@ -328,6 +329,22 @@ class TestClipToBudget:
@pytest.mark.asyncio
async def _fetch_and_expand(repo, document_id, results, max_chars):
"""Do the fetching `expand_context` does, for tests exercising the
expansion logic rather than the batched fetch."""
positions = (
await repo.resolve_refs_grouped(
{document_id: [ref for r in results for ref in r.doc_item_refs]}
)
).get(document_id, {})
items: list = []
if positions:
items = (
await repo.get_items_in_ranges({document_id: window_for(positions)})
).get(document_id, [])
return expand_with_items(results, max_chars, positions, items)
class TestExpandWithItems:
async def test_unresolvable_refs_returns_original(self, temp_db_path):
from haiku.rag.client import HaikuRAG
@ -349,7 +366,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/999999"],
)
assert doc.id is not None
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, doc.id, [result], 5000
)
assert len(expanded) == 1
@ -399,7 +416,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -480,7 +497,7 @@ class TestExpandWithItems:
doc_item_refs=["#/pictures/0"],
page_numbers=[13],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -542,7 +559,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/1", "#/texts/2", "#/texts/3", "#/texts/4"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -574,7 +591,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/tables/0"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 10_000
)
assert len(expanded) == 1
@ -612,7 +629,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/0"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 10_000
)
assert len(expanded) == 1
@ -655,7 +672,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -690,7 +707,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/0"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -720,7 +737,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/2"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -756,7 +773,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/3"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r1, r2], 5000
)
# Ranges around positions 1 and 3 overlap → one merged result.
@ -791,7 +808,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1"],
document_meta={"source_url": "https://example.org/report/view"},
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r1], 5000
)
assert len(expanded) == 1
@ -832,7 +849,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/3"],
chunk_meta={"para_no": "14"},
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r1, r2], 5000
)
assert len(expanded) == 1
@ -873,7 +890,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/3"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r_early, r_best], 5000
)
assert len(expanded) == 1
@ -934,7 +951,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/2"],
page_numbers=[3],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r_low, r_high], 500
)
# The clip window around HIGHMARK cannot contain LOWMARK's item,
@ -990,7 +1007,7 @@ class TestExpandWithItems:
page_numbers=[2],
),
]
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", inputs, 100
)
assert len(expanded) == 2
@ -1037,7 +1054,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/texts/5"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r1, r2], 400
)
assert len(expanded) == 1
@ -1089,7 +1106,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1"],
page_numbers=[2],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r1, r2], 5000
)
assert len(expanded) == 2
@ -1143,7 +1160,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1"],
page_numbers=[8],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository,
"doc-1",
[r_missing_item_page, r_with_item_page],
@ -1190,7 +1207,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/0", "#/texts/1"],
page_numbers=[1, 2],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 100
)
@ -1226,7 +1243,7 @@ class TestExpandWithItems:
document_id="doc-1",
doc_item_refs=["#/tables/0"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 10_000
)
assert len(expanded) == 1
@ -1289,7 +1306,7 @@ class TestExpandWithItemsPictureBytes:
image_data={"#/pictures/0": "BASE64BYTES"},
picture_captions={"#/pictures/0": "Figure 1 caption."},
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
@ -1328,7 +1345,7 @@ class TestExpandWithItemsPictureBytes:
doc_item_refs=["#/pictures/3"],
image_data={"#/pictures/3": "B"},
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r1, r2], 5000
)
# Ranges around positions 1 and 3 overlap → one merged result.
@ -1389,7 +1406,7 @@ class TestExpandWithItemsPictureBytes:
doc_item_refs=["#/pictures/1"],
image_data={"#/pictures/1": "HIGHBYTES"},
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [r_low, r_high], 500
)
assert len(expanded) == 2
@ -1445,10 +1462,10 @@ class TestExpandWithItemsWindowEdges:
)
async def no_window(*_args, **_kwargs):
return []
return {}
monkeypatch.setattr(
rag.document_item_repository, "get_items_in_range", no_window
rag.document_item_repository, "get_items_in_ranges", no_window
)
result = SearchResult(
@ -1457,7 +1474,7 @@ class TestExpandWithItemsWindowEdges:
document_id=doc.id,
doc_item_refs=["#/texts/0"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, doc.id, [result], 5000
)
@ -1502,7 +1519,7 @@ class TestExpandWithItemsWindowEdges:
doc_item_refs=["#/texts/404"],
)
expanded = await expand_with_items(
expanded = await _fetch_and_expand(
rag.document_item_repository, doc.id, [resolvable, unmatched], 5000
)

4
tests/test_context_enhancement.py
View file

@ -322,7 +322,9 @@ async def test_expand_context_single_item_document(temp_db_path):
assert doc.id is not None
# Create a search result with a doc_item_ref pointing to the item
items = await client.document_item_repository.get_items_in_range(doc.id, 0, 10)
items = (
await client.document_item_repository.get_items_in_ranges({doc.id: (0, 10)})
).get(doc.id, [])
assert len(items) > 0
search_results = [

295
tests/test_enrichment_batching.py Normal file
View file

@ -0,0 +1,295 @@
import lancedb
import lancedb.query
import pytest
from haiku.rag.client import HaikuRAG
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:
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_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}
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
async def _seed_expandable(rag: HaikuRAG, document_ids: list[str]) -> None:
"""A section header and two text items, so expansion has something to widen
into. Positions and self_refs repeat across documents."""
for document_id in document_ids:
await rag.document_item_repository.create_items(
document_id,
[
DocumentItem(
document_id=document_id,
position=0,
self_ref="#/texts/0",
label="section_header",
text=f"Section of {document_id}",
),
DocumentItem(
document_id=document_id,
position=1,
self_ref="#/texts/1",
label="text",
text=f"anchor body of {document_id}",
),
DocumentItem(
document_id=document_id,
position=2,
self_ref="#/texts/2",
label="text",
text=f"neighbouring body of {document_id}",
),
],
)
def _text_result(document_id: str) -> SearchResult:
return SearchResult(
chunk_id=f"{document_id}-anchor",
document_id=document_id,
content=f"anchor body of {document_id}",
score=0.9,
doc_item_refs=["#/texts/1"],
)
@pytest.mark.asyncio
async def test_expansion_query_count_is_flat_in_document_count(
temp_db_path, item_queries
):
async with HaikuRAG(temp_db_path, create=True) as rag:
await _seed_expandable(rag, [f"doc-{i}" for i in range(5)])
counts = []
for n in (1, 5):
results = [_text_result(f"doc-{i}") for i in range(n)]
item_queries["n"] = 0
expanded = await rag.expand_context(results)
counts.append(item_queries["n"])
assert len(expanded) == n
assert counts == [2, 2], counts
@pytest.mark.asyncio
async def test_expansion_widens_each_document_with_its_own_items(temp_db_path):
"""Positions repeat across documents, so a batched window fetch keyed on
position alone would splice one document's text into another's context."""
async with HaikuRAG(temp_db_path, create=True) as rag:
await _seed_expandable(rag, ["doc-a", "doc-b"])
expanded = await rag.expand_context(
[_text_result("doc-a"), _text_result("doc-b")]
)
by_doc = {r.document_id: r.content for r in expanded}
assert "neighbouring body of doc-a" in by_doc["doc-a"]
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, item_projections
):
"""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
# 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
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"
@pytest.mark.asyncio
async def test_expansion_keeps_document_order_for_tied_scores(temp_db_path):
"""The score sort is stable, so equal-scored results must come back in the
order they arrived, whether or not their document expands."""
async with HaikuRAG(temp_db_path, create=True) as rag:
await _seed_expandable(rag, ["doc-expandable"])
passthrough = SearchResult(
chunk_id="doc-plain-anchor",
document_id="doc-plain",
content="plain body",
score=0.5,
doc_item_refs=[],
)
expandable = _text_result("doc-expandable")
expandable.score = 0.5
for order in ([passthrough, expandable], [expandable, passthrough]):
expanded = await rag.expand_context(list(order))
assert [r.chunk_id for r in expanded] == [r.chunk_id for r in order]

37
tests/test_picture_in_context.py
View file

@ -16,10 +16,10 @@ 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.context import expand_with_items
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
def _make_png(color: str = "red", size: tuple[int, int] = (4, 4)) -> bytes:
@ -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
@ -205,11 +219,8 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat
doc_item_refs=["#/texts/1"],
labels=["paragraph"],
)
expanded = await expand_with_items(
rag.document_item_repository,
"doc-1",
[seed],
max_chars=10_000,
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [seed], 10_000
)
assert len(expanded) == 1
out = expanded[0]