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

Batch context expansion across documents

Browse source 
`expand_with_items` fetched its own inputs per document: one query to resolve
refs to positions, one for the window of items around them. A result set spanning
N documents cost 2N queries, which was 10 of the 18 measured for a limit=5 search
on a remote object-store corpus.

`expand_context` now does both fetches once for every document it is expanding,
and `expand_with_items` takes the positions and items it needs. Two queries for
one document, and two for five.

Each document keeps its own inclusive window in `get_items_in_ranges`. Positions
repeat across documents, so a shared range would splice one document's items into
another's context.
This commit is contained in:
Yiorgis Gozadinos 2026-08-18 16:35:17 +03:00
parent af6a6b0bbe
commit 5b0444043a
No known key found for this signature in database
7 changed files with 222 additions and 70 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
CHANGELOG.md
View file

@ -11,7 +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.
- 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`.
- 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.

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

@ -215,7 +215,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
@ -228,24 +228,38 @@ 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)
}
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
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)
expanded = await expand_with_items(
client.document_item_repository,
doc_id,
doc_results,
max_chars,
for doc_id, doc_results in expandable.items():
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)

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

@ -167,6 +167,58 @@ class DocumentItemRepository:
)
return {row["self_ref"]: row["position"] for row in rows}
async def resolve_refs_grouped(
self, refs_by_document: "Mapping[str, Sequence[str]]"
) -> dict[str, dict[str, int]]:
"""`resolve_refs` 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_items_in_range` across documents in one query.
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.
"""
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)
grouped.setdefault(item.document_id, []).append(item)
for items in grouped.values():
items.sort(key=lambda x: x.position)
return grouped
async def get_item_count(self, document_id: str) -> int:
"""Count items for a document."""
safe_id = escape_sql_string(document_id)

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

78
tests/test_enrichment_batching.py
View file

@ -132,3 +132,81 @@ async def test_caption_ranked_results_take_at_most_four_queries(
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"]

9
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.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:
@ -205,11 +205,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]