From 0bcf34363aa9dee0b1e5d77b0a137a800184e962 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 8 Jul 2026 15:07:36 +0300 Subject: [PATCH] Carry merged chunk ids on SearchResult and Citation --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/context.py | 6 ++ .../haiku/rag/store/models/chunk.py | 6 ++ .../haiku/rag/store/models/citation.py | 7 ++ tests/store/test_citation.py | 41 +++++++++++ tests/test_context.py | 70 +++++++++++++++++++ 6 files changed, 131 insertions(+) create mode 100644 tests/store/test_citation.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f328381..516db165 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index fbea5c85..c7bab5f8 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index d6d3ba6c..7eb1e57b 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/store/models/citation.py b/haiku_rag_slim/haiku/rag/store/models/citation.py index 5a10c7b0..94e53e27 100644 --- a/haiku_rag_slim/haiku/rag/store/models/citation.py +++ b/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, diff --git a/tests/store/test_citation.py b/tests/store/test_citation.py new file mode 100644 index 00000000..1b1d89db --- /dev/null +++ b/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 diff --git a/tests/test_context.py b/tests/test_context.py index 0bbcef3c..3533110a 100644 --- a/tests/test_context.py +++ b/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."""