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

Merge pull request #497 from ggozad/fix/context-expansion

Browse source 
Context expansion: never drop a retrieved result when merging and clipping
This commit is contained in:
Yiorgis Gozadinos 2026-07-14 10:58:07 +03:00 committed by GitHub
commit 8f085101fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 457 additions and 140 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -1,6 +1,11 @@
# Changelog
## [Unreleased]
### Fixed
- Context expansion no longer drops a retrieved result from a merged group when clipping to `search.max_context_chars`; groups whose clip would evict a constituent's evidence are returned as separate expanded results.
- Context expansion now fills missing page metadata from input search results when the referenced items survive in the expanded result.
## [0.65.1] - 2026-07-10
### Added

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

@ -14,9 +14,11 @@ For STRUCTURED documents (containing section_header or title labels):
This lets small sections (e.g., title+authors) grow into neighboring
content. Picture and table matches are exempt: they return their
enclosing section as-is, never crossing section boundaries.
6. Merge overlapping ranges from multiple results in the same document.
Adjacent but non-overlapping ranges stay separate to preserve section
independence.
6. Merge overlapping ranges from multiple results in the same document,
but only when every constituent's matched evidence survives in the
final budget-clipped window; otherwise the group is split back into
per-result windows so no retrieved result is dropped. Adjacent but
non-overlapping ranges stay separate to preserve section independence.
For UNSTRUCTURED documents (no section headers):
Expand outward item-by-item from the match center until the character
@ -269,6 +271,159 @@ def _find_expansion_range(
return _expand_outward(items, center_idx, max_chars, skip_noise=True)
def _group_lost_constituent(built: SearchResult, group: list[SearchResult]) -> bool:
"""Whether any constituent's matched refs were entirely evicted from ``built``.
Fires when the budget clip (or the fragmented-content fallback) left a
constituent with none of its own items in the built result — its evidence
would be silently dropped if the group stayed merged.
"""
surviving = set(built.doc_item_refs)
return any(
r.doc_item_refs and not surviving.intersection(r.doc_item_refs) for r in group
)
def _add_input_pages_for_surviving_refs(
pages: set[int], refs: list[str], original_results: list[SearchResult]
) -> None:
"""Fill missing item-table pages from inputs whose own refs all survived."""
surviving = set(refs)
if not surviving:
return
for result in original_results:
if not result.page_numbers or not result.doc_item_refs:
continue
if set(result.doc_item_refs) <= surviving:
pages.update(result.page_numbers)
def _build_result(
range_start: int,
range_end: int,
original_results: list[SearchResult],
pos_to_item: dict[int, DocumentItem],
has_sections: bool,
max_chars: int,
) -> SearchResult:
"""Build one expanded result from the items in ``[range_start, range_end]``."""
content_parts: list[str] = []
# Char span of each contributing item within the joined content, so
# metadata can be narrowed to whatever survives a budget clip.
item_spans: list[tuple[int, int, DocumentItem]] = []
cursor = 0
separator = "\n\n"
for pos in range(range_start, range_end + 1):
item = pos_to_item.get(pos)
if item is None:
continue
if has_sections and item.label in _NOISE_LABELS:
continue
if item.text:
if content_parts:
cursor += len(separator)
start = cursor
content_parts.append(item.text)
cursor += len(item.text)
item_spans.append((start, cursor, item))
elif item.label == "picture":
# Pictures may legitimately have empty text (no VLM
# description configured). Keep their self_ref so the
# downstream image_data lookup can still attach bytes. They
# occupy a zero-width position in reading order.
item_spans.append((cursor, cursor, item))
all_headings: list[str] = []
for r in original_results:
if r.headings:
all_headings.extend(h for h in r.headings if h not in all_headings)
# Anchor identity (chunk_id, content/refs fallbacks) on the
# best-scoring constituent — the chunk that earned the result its
# rank — rather than whichever sits earliest in the document.
first = max(original_results, key=lambda r: r.score)
chunk_ids: list[str] = []
for r in original_results:
if r.chunk_id and r.chunk_id not in chunk_ids:
chunk_ids.append(r.chunk_id)
joined = separator.join(content_parts)
# Expansion should never return less content than the original chunk.
# This can happen when item texts are fragmented (e.g., docling splits
# formatted HTML list items into many small text nodes); fall back to
# the chunk's own content, described by the chunk's own metadata.
if len(joined) < len(first.content):
base_content, base_spans = first.content, None
else:
base_content, base_spans = joined, item_spans
if len(base_content) > max_chars:
# Clip to the budget, anchored on the primary (highest-scoring)
# chunk, and narrow metadata to whatever survives the window.
win_start, win_end = _clip_window(
base_content, [first, *original_results], max_chars
)
expanded_content = base_content[win_start:win_end]
if base_spans is None:
pages, refs, labels = (
set(first.page_numbers),
list(first.doc_item_refs),
set(first.labels),
)
else:
pages, refs, labels = _collect_meta(
[s for s in base_spans if _span_in_window(s, win_start, win_end)]
)
else:
expanded_content = base_content
if base_spans is None:
pages, refs, labels = (
set(first.page_numbers),
list(first.doc_item_refs),
set(first.labels),
)
else:
pages, refs, labels = _collect_meta(base_spans)
_add_input_pages_for_surviving_refs(pages, refs, original_results)
# Carry image_data and picture_captions from the originally retrieved
# chunks, but only for constituents whose refs survive the window — a
# chunk clipped out of the budget must not still ship its image to the
# model. Pictures swept in by section expansion are referenced in
# ``refs`` but their bytes are never re-fetched, so the multimodal
# payload stays bounded to what was actually retrieved and shown.
surviving_refs = set(refs)
merged_image_data: dict[str, str] = {}
merged_captions: dict[str, str] = {}
for r in original_results:
if r.doc_item_refs and not surviving_refs.intersection(r.doc_item_refs):
continue
if r.image_data:
merged_image_data.update(r.image_data)
if r.picture_captions:
merged_captions.update(r.picture_captions)
return SearchResult(
content=expanded_content,
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
chunk_ids=chunk_ids,
document_id=first.document_id,
document_uri=first.document_uri,
document_title=first.document_title,
doc_item_refs=refs or first.doc_item_refs,
page_numbers=sorted(pages) or first.page_numbers,
headings=all_headings or None,
labels=sorted(labels) or first.labels,
image_data=merged_image_data or None,
picture_captions=merged_captions,
)
_WINDOW_MARGIN = 100
@ -317,125 +472,27 @@ async def expand_with_items(
ranges.append((lo, hi, result))
merged = _merge_ranges(ranges)
constituent_range = {id(result): (lo, hi) for lo, hi, result in ranges}
# Build results from the window items
pos_to_item = {item.position: item for item in window_items}
final_results: list[SearchResult] = []
for range_start, range_end, original_results in merged:
content_parts: list[str] = []
# Char span of each contributing item within the joined content, so
# metadata can be narrowed to whatever survives a budget clip.
item_spans: list[tuple[int, int, DocumentItem]] = []
cursor = 0
separator = "\n\n"
for pos in range(range_start, range_end + 1):
item = pos_to_item.get(pos)
if item is None:
continue
if has_sections and item.label in _NOISE_LABELS:
continue
if item.text:
if content_parts:
cursor += len(separator)
start = cursor
content_parts.append(item.text)
cursor += len(item.text)
item_spans.append((start, cursor, item))
elif item.label == "picture":
# Pictures may legitimately have empty text (no VLM
# description configured). Keep their self_ref so the
# downstream image_data lookup can still attach bytes. They
# occupy a zero-width position in reading order.
item_spans.append((cursor, cursor, item))
all_headings: list[str] = []
for r in original_results:
if r.headings:
all_headings.extend(h for h in r.headings if h not in all_headings)
# Anchor identity (chunk_id, content/refs fallbacks) on the
# best-scoring constituent — the chunk that earned the result its
# rank — rather than whichever sits earliest in the document.
first = max(original_results, key=lambda r: r.score)
chunk_ids: list[str] = []
for r in original_results:
if r.chunk_id and r.chunk_id not in chunk_ids:
chunk_ids.append(r.chunk_id)
joined = separator.join(content_parts)
# Expansion should never return less content than the original chunk.
# This can happen when item texts are fragmented (e.g., docling splits
# formatted HTML list items into many small text nodes); fall back to
# the chunk's own content, described by the chunk's own metadata.
if len(joined) < len(first.content):
base_content, base_spans = first.content, None
else:
base_content, base_spans = joined, item_spans
if len(base_content) > max_chars:
# Clip to the budget, anchored on the primary (highest-scoring)
# chunk, and narrow metadata to whatever survives the window.
win_start, win_end = _clip_window(
base_content, [first, *original_results], max_chars
)
expanded_content = base_content[win_start:win_end]
if base_spans is None:
pages, refs, labels = (
set(first.page_numbers),
list(first.doc_item_refs),
set(first.labels),
)
else:
pages, refs, labels = _collect_meta(
[s for s in base_spans if _span_in_window(s, win_start, win_end)]
)
else:
expanded_content = base_content
if base_spans is None:
pages, refs, labels = (
set(first.page_numbers),
list(first.doc_item_refs),
set(first.labels),
)
else:
pages, refs, labels = _collect_meta(base_spans)
# Carry image_data and picture_captions from the originally retrieved
# chunks, but only for constituents whose refs survive the window — a
# chunk clipped out of the budget must not still ship its image to the
# model. Pictures swept in by section expansion are referenced in
# ``refs`` but their bytes are never re-fetched, so the multimodal
# payload stays bounded to what was actually retrieved and shown.
surviving_refs = set(refs)
merged_image_data: dict[str, str] = {}
merged_captions: dict[str, str] = {}
for r in original_results:
if r.doc_item_refs and not surviving_refs.intersection(r.doc_item_refs):
continue
if r.image_data:
merged_image_data.update(r.image_data)
if r.picture_captions:
merged_captions.update(r.picture_captions)
final_results.append(
SearchResult(
content=expanded_content,
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
chunk_ids=chunk_ids,
document_id=first.document_id,
document_uri=first.document_uri,
document_title=first.document_title,
doc_item_refs=refs or first.doc_item_refs,
page_numbers=sorted(pages) or first.page_numbers,
headings=all_headings or None,
labels=sorted(labels) or first.labels,
image_data=merged_image_data or None,
picture_captions=merged_captions,
)
for range_start, range_end, group in merged:
built = _build_result(
range_start, range_end, group, pos_to_item, has_sections, max_chars
)
if len(group) > 1 and _group_lost_constituent(built, group):
# The merged window cannot afford every constituent's evidence:
# un-merge so no retrieved result is dropped from the group. Each
# constituent gets its own window, clipped around its own anchor.
for result in group:
lo, hi = constituent_range[id(result)]
final_results.append(
_build_result(
lo, hi, [result], pos_to_item, has_sections, max_chars
)
)
continue
final_results.append(built)
return final_results + passthrough

299
tests/test_context.py
View file

@ -810,12 +810,11 @@ class TestExpandWithItems:
# provenance still lists both
assert set(expanded[0].chunk_ids) == {"c-early", "c-best"}
async def test_clipped_merged_result_keeps_anchor_evidence_and_pages(
self, temp_db_path
):
"""When a merged result is clipped to budget, the surviving window is
centered on the anchor (highest-scoring) chunk, and page_numbers reflect
only the content that survived — not the full merged range."""
async def test_clipped_merge_that_evicts_a_constituent_splits(self, temp_db_path):
"""A merged group whose budget clip would evict a constituent's evidence
is split back into per-result windows: no retrieved result is dropped.
Each split result keeps its own evidence and metadata describing only
its visible content."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
@ -866,17 +865,268 @@ class TestExpandWithItems:
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [r_low, r_high], 500
)
# The clip window around HIGHMARK cannot contain LOWMARK's item,
# so the group splits instead of dropping r_low.
assert len(expanded) == 2
by_chunk = {e.chunk_id: e for e in expanded}
e_high = by_chunk["c-high"]
assert "HIGHMARK" in e_high.content
assert "LOWMARK" not in e_high.content
assert 3 in e_high.page_numbers
# per-result metadata still describes only the visible content
assert 1 not in e_high.page_numbers
assert "#/texts/0" not in e_high.doc_item_refs
assert e_high.chunk_ids == ["c-high"]
e_low = by_chunk["c-low"]
assert "LOWMARK" in e_low.content
assert 1 in e_low.page_numbers
assert e_low.chunk_ids == ["c-low"]
async def test_split_results_each_keep_own_evidence_and_budget(self, temp_db_path):
"""Close matches on different pages at a small budget: the merged
window cannot afford both, so each hit gets its own clipped window.
Every input result's page survives across the output results."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=pos,
self_ref=f"#/texts/{pos}",
label="text",
text=f"i{pos:02d}" + "x" * 17,
page_numbers=[1 if pos < 3 else 2],
)
for pos in range(8)
]
await rag.document_item_repository.create_items("doc-1", items)
inputs = [
SearchResult(
content=items[1].text,
score=0.5,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
page_numbers=[1],
),
SearchResult(
content=items[5].text,
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/5"],
page_numbers=[2],
),
]
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", inputs, 100
)
assert len(expanded) == 2
for e in expanded:
assert len(e.content) <= 100
contents = " || ".join(e.content for e in expanded)
assert items[1].text in contents
assert items[5].text in contents
input_pages = {p for r in inputs for p in r.page_numbers}
output_pages = {p for e in expanded for p in e.page_numbers}
assert input_pages <= output_pages
async def test_clipped_merge_stays_merged_when_all_evidence_survives(
self, temp_db_path
):
"""A merged group is clipped but the window still contains every
constituent's evidence: no split, one merged result."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=pos,
self_ref=f"#/texts/{pos}",
label="text",
text=f"i{pos:02d}" + "y" * 97,
)
for pos in range(10)
]
await rag.document_item_repository.create_items("doc-1", items)
r1 = SearchResult(
content=items[4].text,
score=0.5,
chunk_id="c1",
document_id="doc-1",
doc_item_refs=["#/texts/4"],
)
r2 = SearchResult(
content=items[5].text,
score=0.9,
chunk_id="c2",
document_id="doc-1",
doc_item_refs=["#/texts/5"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [r1, r2], 400
)
assert len(expanded) == 1
e = expanded[0]
# anchor is the high-scoring chunk, and its evidence survives clipping
assert e.chunk_id == "c-high"
assert "HIGHMARK" in e.content
assert "LOWMARK" not in e.content
# page_numbers reflect only the surviving window, not the full range
assert 3 in e.page_numbers
assert 1 not in e.page_numbers
# refs likewise exclude the clipped-out item
assert "#/texts/0" not in e.doc_item_refs
assert len(e.content) <= 400
assert items[4].text in e.content
assert items[5].text in e.content
assert set(e.chunk_ids) == {"c1", "c2"}
async def test_fragmented_merge_splits_instead_of_dropping(self, temp_db_path):
"""When fragmented item text triggers the chunk-content fallback for a
merged group, the group splits so the non-primary hit is not dropped."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="text",
text="frag a",
page_numbers=[1],
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="frag b",
page_numbers=[2],
),
]
await rag.document_item_repository.create_items("doc-1", items)
r1 = SearchResult(
content="A" * 500,
score=0.9,
chunk_id="c1",
document_id="doc-1",
doc_item_refs=["#/texts/0"],
page_numbers=[1],
)
r2 = SearchResult(
content="B" * 400,
score=0.5,
chunk_id="c2",
document_id="doc-1",
doc_item_refs=["#/texts/1"],
page_numbers=[2],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [r1, r2], 5000
)
assert len(expanded) == 2
by_chunk = {e.chunk_id: e for e in expanded}
assert by_chunk["c1"].content == "A" * 500
assert by_chunk["c1"].page_numbers == [1]
assert by_chunk["c2"].content == "B" * 400
assert by_chunk["c2"].page_numbers == [2]
async def test_surviving_refs_fill_missing_item_pages_from_input(
self, temp_db_path
):
"""When a visible item has missing page metadata, use the input
result's page_numbers as a floor for that surviving ref."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="text",
text="Visible item with missing item-table pages.",
page_numbers=[],
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Visible item with stored item-table pages.",
page_numbers=[8],
),
]
await rag.document_item_repository.create_items("doc-1", items)
r_missing_item_page = SearchResult(
content=items[0].text,
score=0.9,
chunk_id="c-missing",
document_id="doc-1",
doc_item_refs=["#/texts/0"],
page_numbers=[7],
)
r_with_item_page = SearchResult(
content=items[1].text,
score=0.8,
chunk_id="c-present",
document_id="doc-1",
doc_item_refs=["#/texts/1"],
page_numbers=[8],
)
expanded = await expand_with_items(
rag.document_item_repository,
"doc-1",
[r_missing_item_page, r_with_item_page],
5000,
)
assert len(expanded) == 1
assert set(expanded[0].doc_item_refs) == {"#/texts/0", "#/texts/1"}
assert expanded[0].page_numbers == [7, 8]
async def test_input_pages_not_added_for_clipped_out_refs(self, temp_db_path):
"""Input page metadata is not blindly unioned when only some of a
constituent's refs survive the clip window."""
from haiku.rag.client import HaikuRAG
item0_text = "LEFTMARK " + "a" * 91
item1_text = "RIGHTMARK " + "b" * 70
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="text",
text=item0_text,
page_numbers=[1],
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text=item1_text,
page_numbers=[2],
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content=item0_text,
score=0.9,
chunk_id="c-both",
document_id="doc-1",
doc_item_refs=["#/texts/0", "#/texts/1"],
page_numbers=[1, 2],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 100
)
assert len(expanded) == 1
assert "LEFTMARK" in expanded[0].content
assert "RIGHTMARK" not in expanded[0].content
assert expanded[0].doc_item_refs == ["#/texts/0"]
assert expanded[0].page_numbers == [1]
async def test_fuzzy_match_preserves_central_marker(self, temp_db_path):
"""The chunk's text need not be verbatim in the joined item text: a clean
@ -1016,10 +1266,10 @@ class TestExpandWithItemsPictureBytes:
"#/pictures/3": "B",
}
async def test_clipped_out_picture_bytes_dropped(self, temp_db_path):
"""A lower-scoring picture chunk clipped out of the budget window no
longer contributes its image bytes — the model must not receive an
image the citation and visualization omit."""
async def test_split_results_carry_only_own_picture_bytes(self, temp_db_path):
"""When a clipped merge splits, each result ships only the image bytes
its own window shows — the model must not receive an image the
citation and visualization omit."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
@ -1070,6 +1320,11 @@ class TestExpandWithItemsPictureBytes:
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [r_low, r_high], 500
)
assert len(expanded) == 1
assert "#/pictures/0" not in expanded[0].doc_item_refs
assert expanded[0].image_data == {"#/pictures/1": "HIGHBYTES"}
assert len(expanded) == 2
by_chunk = {e.chunk_id: e for e in expanded}
e_high = by_chunk["c-high"]
assert "#/pictures/0" not in e_high.doc_item_refs
assert e_high.image_data == {"#/pictures/1": "HIGHBYTES"}
e_low = by_chunk["c-low"]
assert e_low.image_data == {"#/pictures/0": "LOWBYTES"}
assert "HIGHBYTES" not in (e_low.image_data or {}).values()