Carry merged chunk ids on SearchResult and Citation

CHANGELOG.md

@@ -16,6 +16,7 @@
 ### Added
 
 - `update_document` accepts a `uri` argument to change a document's URI.
+- `SearchResult.chunk_ids` and `Citation.chunk_ids` carry the chunk ids merged into an expanded result.
 - docling-serve requests fail over to another instance on transport/5xx errors and skip instances whose circuit breaker is open; tune via `providers.docling_serve.max_attempts` and `providers.docling_serve.circuit_breaker`.
 
 ### Fixed

haiku_rag_slim/haiku/rag/context.py

@@ -335,6 +335,11 @@ async def expand_with_items(
 
         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)
+
         # 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).
@@ -351,6 +356,7 @@ async def expand_with_items(
                 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,

haiku_rag_slim/haiku/rag/store/models/chunk.py

@@ -124,11 +124,17 @@ class SearchResult(BaseModel):
     when the caller asked to omit them via ``include_images=False`` on
     ``client.search``. Same shape is used everywhere — MCP, in-process search,
     agent toolsets — so non-vision callers see ``None`` and pay nothing.
+
+    ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
+    into this result; empty means just ``chunk_id``. It lets citation
+    consumers (visual grounding) reproduce a merged expansion and is never
+    part of ``format_for_agent`` output.
     """
 
     content: str
     score: float
     chunk_id: str | None = None
+    chunk_ids: list[str] = []
     document_id: str | None = None
     document_uri: str | None = None
     document_title: str | None = None

haiku_rag_slim/haiku/rag/store/models/citation.py

@@ -18,11 +18,17 @@ class Citation(BaseModel):
     cited chunk. Empty for text-only citations. UIs can fetch the picture
     bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
     and render them alongside the text content.
+
+    ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
+    into the cited result (always includes ``chunk_id``). Visual grounding
+    passes them all to ``visualize_chunk`` so the rendered pages reproduce
+    the merged expansion.
     """
 
     index: int | None = None
     document_id: str
     chunk_id: str
+    chunk_ids: list[str] = Field(default_factory=list)
     document_uri: str
     document_title: str | None = None
     page_numbers: list[int] = Field(default_factory=list)
@@ -51,6 +57,7 @@ def resolve_citations(
             Citation(
                 document_id=r.document_id or "",
                 chunk_id=chunk_id,
+                chunk_ids=r.chunk_ids or [chunk_id],
                 document_uri=r.document_uri or "",
                 document_title=r.document_title,
                 page_numbers=r.page_numbers,

tests/store/test_citation.py

@@ -0,0 +1,41 @@
+from haiku.rag.store.models.chunk import SearchResult
+from haiku.rag.store.models.citation import resolve_citations
+
+
+def _result(chunk_id: str, chunk_ids: list[str] | None = None) -> SearchResult:
+    return SearchResult(
+        content="content",
+        score=0.9,
+        chunk_id=chunk_id,
+        chunk_ids=chunk_ids or [],
+        document_id="doc-1",
+        document_uri="test://doc",
+    )
+
+
+def test_resolve_citations_copies_merged_chunk_ids():
+    result = _result("c1", chunk_ids=["c1", "c2"])
+    citations = resolve_citations(["c1"], [result])
+    assert len(citations) == 1
+    assert citations[0].chunk_id == "c1"
+    assert citations[0].chunk_ids == ["c1", "c2"]
+
+
+def test_resolve_citations_falls_back_to_chunk_id():
+    result = _result("c1")
+    citations = resolve_citations(["c1"], [result])
+    assert len(citations) == 1
+    assert citations[0].chunk_ids == ["c1"]
+
+
+def test_resolve_citations_strips_brackets():
+    result = _result("c1")
+    citations = resolve_citations(["[c1]"], [result])
+    assert len(citations) == 1
+    assert citations[0].chunk_id == "c1"
+
+
+def test_resolve_citations_skips_unknown_ids():
+    result = _result("c1")
+    citations = resolve_citations(["c1", "missing"], [result])
+    assert len(citations) == 1

tests/test_context.py

@@ -697,6 +697,76 @@ class TestExpandWithItems:
             assert len(expanded[0].content) <= 5000
             assert marker in expanded[0].content
 
+    async def test_solo_result_carries_own_chunk_id(self, temp_db_path):
+        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)
+
+            result = SearchResult(
+                content="Paragraph 2.",
+                score=0.9,
+                chunk_id="c1",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/2"],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [result], 5000
+            )
+            assert len(expanded) == 1
+            assert expanded[0].chunk_ids == ["c1"]
+
+    async def test_merged_results_carry_all_chunk_ids(self, temp_db_path):
+        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)
+
+            r1 = SearchResult(
+                content="Paragraph 1.",
+                score=0.9,
+                chunk_id="c1",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/1"],
+            )
+            r2 = SearchResult(
+                content="Paragraph 3.",
+                score=0.85,
+                chunk_id="c2",
+                document_id="doc-1",
+                doc_item_refs=["#/texts/3"],
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [r1, r2], 5000
+            )
+            # Ranges around positions 1 and 3 overlap → one merged result.
+            assert len(expanded) == 1
+            assert expanded[0].chunk_id == "c1"
+            assert expanded[0].chunk_ids == ["c1", "c2"]
+            # Sibling chunk ids are plumbing for visualization, never shown
+            # to the model.
+            assert "c2" not in expanded[0].format_for_agent()
+
     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."""