Anchor and clip merged citations on the highest-scoring chunk

A merged search result took its chunk_id from whichever constituent
sorted earliest in the document, while its score was the max across the
group — so the citation's identity could point at a different, less
relevant chunk. Anchor chunk_id and the content/refs fallbacks on the
max-score constituent, clip the budget window around that same chunk so
its evidence is never trimmed away, and narrow page_numbers, doc_item_refs,
and attached image bytes to the items that survive the clip.

haiku_rag_slim/haiku/rag/context.py

@@ -104,16 +104,18 @@ def _evidence_anchors(content: str, max_chars: int) -> list[str]:
     return anchors
 
 
-def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
+def _clip_window(
-    """Clip expanded content to ``max_chars``, keeping the matched evidence.
+    content: str, results: list[SearchResult], max_chars: int
+) -> tuple[int, int]:
+    """Return the ``[start, end)`` char window to keep when clipping ``content``.
 
     Anchors on the first locatable result in ``results`` order (the primary
     chunk that supplies the expanded result's identity) and returns a
-    ``max_chars``-wide window centered on it. Falls back to a prefix cut only
+    ``max_chars``-wide window centered on it. Falls back to a prefix window only
     when no anchor is locatable (heavy drift).
     """
     if max_chars <= 0:
-        return ""
+        return (0, 0)
     evidence_start, evidence_len = -1, 0
     for result in results:
         for anchor in _evidence_anchors(result.content, max_chars):
@@ -124,14 +126,45 @@ def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -
         if evidence_start != -1:
             break
     if evidence_start == -1:
-        return content[:max_chars]
+        return (0, min(len(content), max_chars))
     center = evidence_start + evidence_len // 2
     start = max(0, center - max_chars // 2)
     end = min(len(content), start + max_chars)
     start = max(0, end - max_chars)
+    return (start, end)
+
+
+def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
+    """Clip expanded content to ``max_chars``, keeping the matched evidence."""
+    start, end = _clip_window(content, results, max_chars)
     return content[start:end]
 
 
+def _collect_meta(
+    spans: list[tuple[int, int, DocumentItem]],
+) -> tuple[set[int], list[str], set[str]]:
+    """Union the page numbers, refs, and labels of the given item spans."""
+    pages: set[int] = set()
+    refs: list[str] = []
+    labels: set[str] = set()
+    for _start, _end, item in spans:
+        refs.append(item.self_ref)
+        if item.label:
+            labels.add(item.label)
+        pages.update(item.page_numbers)
+    return pages, refs, labels
+
+
+def _span_in_window(
+    span: tuple[int, int, DocumentItem], win_start: int, win_end: int
+) -> bool:
+    """Whether an item's char span overlaps the ``[win_start, win_end]`` clip window."""
+    start, end, _item = span
+    if start == end:  # zero-width picture position
+        return win_start <= start <= win_end
+    return start < win_end and end > win_start
+
+
 def _expand_outward(
     items: list[DocumentItem],
     center_idx: int,
@@ -290,9 +323,11 @@ async def expand_with_items(
     final_results: list[SearchResult] = []
     for range_start, range_end, original_results in merged:
         content_parts: list[str] = []
-        refs: list[str] = []
+        # Char span of each contributing item within the joined content, so
-        pages: set[int] = set()
+        # metadata can be narrowed to whatever survives a budget clip.
-        labels: set[str] = set()
+        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)
@@ -301,55 +336,89 @@ async def expand_with_items(
             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)
-                refs.append(item.self_ref)
+                cursor += len(item.text)
-                if item.label:
+                item_spans.append((start, cursor, item))
-                    labels.add(item.label)
-                pages.update(item.page_numbers)
             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.
+                # downstream image_data lookup can still attach bytes. They
-                refs.append(item.self_ref)
+                # occupy a zero-width position in reading order.
-                labels.add(item.label)
+                item_spans.append((cursor, cursor, item))
-                pages.update(item.page_numbers)
 
         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)
 
-        # Carry image_data and picture_captions through expansion so that
+        # Anchor identity (chunk_id, content/refs fallbacks) on the
-        # only pictures from the originally retrieved chunks get attached.
+        # best-scoring constituent — the chunk that earned the result its
-        # Pictures swept in by section expansion are referenced in `refs`
+        # rank — rather than whichever sits earliest in the document.
-        # for cross-referencing but their bytes are not re-fetched —
+        first = max(original_results, key=lambda r: r.score)
-        # otherwise a single search can balloon the response with adjacent
-        # figures the model did not actually retrieve.
-        merged_image_data: dict[str, str] = {}
-        merged_captions: dict[str, str] = {}
-        for r in original_results:
-            if r.image_data:
-                merged_image_data.update(r.image_data)
-            if r.picture_captions:
-                merged_captions.update(r.picture_captions)
-
-        first = original_results[0]
 
         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).
+        # formatted HTML list items into many small text nodes); fall back to
-        expanded_content = "\n\n".join(content_parts)
+        # the chunk's own content, described by the chunk's own metadata.
-        if len(expanded_content) < len(first.content):
+        if len(joined) < len(first.content):
-            expanded_content = first.content
+            base_content, base_spans = first.content, None
-        if len(expanded_content) > max_chars:
+        else:
-            expanded_content = _clip_to_budget(
+            base_content, base_spans = joined, item_spans
-                expanded_content, original_results, max_chars
+
+        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(

tests/test_context.py

@@ -767,6 +767,117 @@ class TestExpandWithItems:
             # to the model.
             assert "c2" not in expanded[0].format_for_agent()
 
+    async def test_merged_anchor_is_highest_scoring_constituent(self, temp_db_path):
+        """A merged result's chunk_id anchors on the best-scoring constituent,
+        not whichever chunk sits earliest in the document."""
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            items = [
+                DocumentItem(
+                    document_id="doc-1",
+                    position=i,
+                    self_ref=f"#/texts/{i}",
+                    label="text",
+                    text=f"Paragraph {i}. " * 10,
+                )
+                for i in range(5)
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            # earlier in the document, lower score
+            r_early = SearchResult(
+                content="Paragraph 1.",
+                score=0.40,
+                chunk_id="c-early",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/1"],
+            )
+            # later in the document, higher score — the real hit
+            r_best = SearchResult(
+                content="Paragraph 3.",
+                score=0.95,
+                chunk_id="c-best",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/3"],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r_early, r_best], 5000
+            )
+            assert len(expanded) == 1
+            assert expanded[0].chunk_id == "c-best"
+            assert expanded[0].score == 0.95
+            # 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."""
+        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="LOWMARK " + "a" * 400,
+                    page_numbers=[1],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="text",
+                    text="b" * 400,
+                    page_numbers=[2],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=2,
+                    self_ref="#/texts/2",
+                    label="text",
+                    text="c" * 400 + " HIGHMARK",
+                    page_numbers=[3],
+                ),
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            r_low = SearchResult(
+                content="LOWMARK " + "a" * 400,
+                score=0.4,
+                chunk_id="c-low",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/0"],
+                page_numbers=[1],
+            )
+            r_high = SearchResult(
+                content="c" * 400 + " HIGHMARK",
+                score=0.9,
+                chunk_id="c-high",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/2"],
+                page_numbers=[3],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r_low, r_high], 500
+            )
+            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
+
     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
         central marker is still located via the central-slice anchor."""
@@ -904,3 +1015,61 @@ class TestExpandWithItemsPictureBytes:
                 "#/pictures/1": "A",
                 "#/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."""
+        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="#/pictures/0",
+                    label="picture",
+                    text="LOWPIC " + "a" * 400,
+                    page_numbers=[1],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=1,
+                    self_ref="#/texts/1",
+                    label="text",
+                    text="b" * 400,
+                    page_numbers=[2],
+                ),
+                DocumentItem(
+                    document_id="doc-1",
+                    position=2,
+                    self_ref="#/pictures/1",
+                    label="picture",
+                    text="c" * 400 + " HIGHPIC",
+                    page_numbers=[3],
+                ),
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            r_low = SearchResult(
+                content="LOWPIC " + "a" * 400,
+                score=0.4,
+                chunk_id="c-low",
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/0"],
+                image_data={"#/pictures/0": "LOWBYTES"},
+            )
+            r_high = SearchResult(
+                content="c" * 400 + " HIGHPIC",
+                score=0.9,
+                chunk_id="c-high",
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/1"],
+                image_data={"#/pictures/1": "HIGHBYTES"},
+            )
+            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"}