diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f8976dd..f7c93ca3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # 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. + ## [0.65.1] - 2026-07-10 ### Added diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 0b6c37ef..c49e66d0 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -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,143 @@ 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 _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) + + # 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 +456,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 diff --git a/tests/test_context.py b/tests/test_context.py index 5bc6dcf3..31f30c2b 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -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,167 @@ 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_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 +1165,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 +1219,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()