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

Context expansion: never drop a retrieved result when merging and clipping

CHANGELOG.md

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

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.
+  6. Merge overlapping ranges from multiple results in the same document,
-     Adjacent but non-overlapping ranges stay separate to preserve section
+     but only when every constituent's matched evidence survives in the
-     independence.
+     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,59 +271,42 @@ def _find_expansion_range(
     return _expand_outward(items, center_idx, max_chars, skip_noise=True)
 
 
-_WINDOW_MARGIN = 100
+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
-async def expand_with_items(
+    constituent with none of its own items in the built result — its evidence
-    document_item_repository: DocumentItemRepository,
+    would be silently dropped if the group stayed merged.
-    document_id: str,
+    """
-    results: list[SearchResult],
+    surviving = set(built.doc_item_refs)
-    max_chars: int,
+    return any(
-) -> list[SearchResult]:
+        r.doc_item_refs and not surviving.intersection(r.doc_item_refs) for r in group
-    """Expand results using the document_items table."""
-    all_refs = []
-    for result in results:
-        all_refs.extend(result.doc_item_refs)
-
-    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:
-        return results
 
-    has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items)
+def _add_input_pages_for_surviving_refs(
-
+    pages: set[int], refs: list[str], original_results: list[SearchResult]
-    # Compute expansion ranges per result
+) -> None:
-    ranges: list[tuple[int, int, SearchResult]] = []
+    """Fill missing item-table pages from inputs whose own refs all survived."""
-    passthrough: list[SearchResult] = []
+    surviving = set(refs)
-
+    if not surviving:
-    for result in results:
+        return
-        matched = {ref_positions[r] for r in result.doc_item_refs if r in ref_positions}
+    for result in original_results:
-        if not matched:
+        if not result.page_numbers or not result.doc_item_refs:
-            passthrough.append(result)
             continue
+        if set(result.doc_item_refs) <= surviving:
+            pages.update(result.page_numbers)
 
-        lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
-        ranges.append((lo, hi, result))
 
-    merged = _merge_ranges(ranges)
+def _build_result(
-
+    range_start: int,
-    # Build results from the window items
+    range_end: int,
-    pos_to_item = {item.position: item for item in window_items}
+    original_results: list[SearchResult],
-    final_results: list[SearchResult] = []
+    pos_to_item: dict[int, DocumentItem],
-    for range_start, range_end, original_results in merged:
+    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.
@@ -403,6 +388,8 @@ async def expand_with_items(
         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
@@ -420,8 +407,7 @@ async def expand_with_items(
         if r.picture_captions:
             merged_captions.update(r.picture_captions)
 
-        final_results.append(
+    return SearchResult(
-            SearchResult(
         content=expanded_content,
         score=max(r.score for r in original_results),
         chunk_id=first.chunk_id,
@@ -436,6 +422,77 @@ async def expand_with_items(
         image_data=merged_image_data or None,
         picture_captions=merged_captions,
     )
+
+
+_WINDOW_MARGIN = 100
+
+
+async def expand_with_items(
+    document_item_repository: DocumentItemRepository,
+    document_id: str,
+    results: list[SearchResult],
+    max_chars: int,
+) -> list[SearchResult]:
+    """Expand results using the document_items table."""
+    all_refs = []
+    for result in results:
+        all_refs.extend(result.doc_item_refs)
+
+    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:
+        return results
+
+    has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items)
+
+    # Compute expansion ranges per result
+    ranges: list[tuple[int, int, SearchResult]] = []
+    passthrough: list[SearchResult] = []
+
+    for result in results:
+        matched = {ref_positions[r] for r in result.doc_item_refs if r in ref_positions}
+        if not matched:
+            passthrough.append(result)
+            continue
+
+        lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
+        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, 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

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(
+    async def test_clipped_merge_that_evicts_a_constituent_splits(self, temp_db_path):
-        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.
-        """When a merged result is clipped to budget, the surviving window is
+        Each split result keeps its own evidence and metadata describing only
-        centered on the anchor (highest-scoring) chunk, and page_numbers reflect
+        its visible content."""
-        only the content that survived — not the full merged range."""
         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 len(e.content) <= 400
-            assert e.chunk_id == "c-high"
+            assert items[4].text in e.content
-            assert "HIGHMARK" in e.content
+            assert items[5].text in e.content
-            assert "LOWMARK" not in e.content
+            assert set(e.chunk_ids) == {"c1", "c2"}
-            # page_numbers reflect only the surviving window, not the full range
+
-            assert 3 in e.page_numbers
+    async def test_fragmented_merge_splits_instead_of_dropping(self, temp_db_path):
-            assert 1 not in e.page_numbers
+        """When fragmented item text triggers the chunk-content fallback for a
-            # refs likewise exclude the clipped-out item
+        merged group, the group splits so the non-primary hit is not dropped."""
-            assert "#/texts/0" not in e.doc_item_refs
+        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):
+    async def test_split_results_carry_only_own_picture_bytes(self, temp_db_path):
-        """A lower-scoring picture chunk clipped out of the budget window no
+        """When a clipped merge splits, each result ships only the image bytes
-        longer contributes its image bytes — the model must not receive an
+        its own window shows — the model must not receive an image the
-        image the citation and visualization omit."""
+        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 len(expanded) == 2
-            assert "#/pictures/0" not in expanded[0].doc_item_refs
+            by_chunk = {e.chunk_id: e for e in expanded}
-            assert expanded[0].image_data == {"#/pictures/1": "HIGHBYTES"}
+            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()