From 225a37ca73b3de1cdae631b8922e4346c176e2b9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Aug 2026 10:35:46 +0300 Subject: [PATCH] Tell one collection's pictures from another's Three places identified a picture by document and reference alone, and one identified a citation's images by chunk id alone. Both repeat between copies of a database, so a search returning a figure from two collections sent one, a capsule retained one, and a citation rendered the other collection's figures. Keyed on the source as well: `(source, document_id, self_ref)` for search pictures, the capability and source for retained ones, and `qualified_id` for the chat's citation images. --- .../haiku/rag/capabilities/compaction.py | 19 +++-- haiku_rag_slim/haiku/rag/chat/app.py | 7 +- .../haiku/rag/chat/widgets/chat_history.py | 11 ++- haiku_rag_slim/haiku/rag/tools/search.py | 8 +- tests/capabilities/test_evidence_capsule.py | 26 +++++++ tests/chat/test_chat_app.py | 73 +++++++++++++++++++ tests/tools/test_search.py | 31 ++++++++ 7 files changed, 159 insertions(+), 16 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index ed171eb0..4d77a378 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -60,9 +60,9 @@ def picture_label(chunk_id: str, self_ref: str) -> str: class RetainedPicture: """A picture to re-attach, with the label that must accompany it. - Addressed by owner, document and reference, because a reference such as - ``#/pictures/0`` repeats across documents and capabilities. The label travels - with it so it can never be emitted without its image. + Addressed by owner, collection, document and reference, because a reference + such as ``#/pictures/0`` repeats across all three. The label travels with it + so it can never be emitted without its image. """ capability: str @@ -161,7 +161,7 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: lines = [CAPSULE_HEADER] pictures: list[RetainedPicture] = [] - seen: set[tuple[str, str, str]] = set() + seen: set[tuple[str, str | None, str, str]] = set() # A capsule may combine citations from different search scopes. include_collection = len({entry.citation.source for entry in entries}) > 1 position = 0 @@ -174,9 +174,14 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: lines.append(entry.render(include_collection=include_collection)) for self_ref in entry.citation.picture_refs: # Overlapping chunks cite one figure, and a provider counts it twice. - # Identity is owner plus document plus reference, so the same reference - # in another document stays a different picture. - identity = (entry.capability, entry.citation.document_id, self_ref) + # A reference such as `#/pictures/0` repeats across documents, and a + # document repeats across collections. + identity = ( + entry.capability, + entry.citation.source, + entry.citation.document_id, + self_ref, + ) if identity in seen: continue seen.add(identity) diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index a9b8a1fb..0d110435 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -36,6 +36,7 @@ from haiku.rag.chat.widgets.prompt import ( ) from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config +from haiku.rag.store.models.chunk import qualified_id from haiku.rag.telemetry import configure as configure_telemetry configure_telemetry(service_name="haiku-rag") @@ -294,7 +295,7 @@ class ChatApp(App): if not citations: return - picture_bytes: dict[str, list[bytes]] = {} + picture_bytes: dict[tuple[str | None, str | None], list[bytes]] = {} if self.client is not None: for citation in citations: refs = list(citation.picture_refs or []) @@ -312,7 +313,9 @@ class ChatApp(App): if data: blobs.append(data) if blobs: - picture_bytes[citation.chunk_id] = blobs + picture_bytes[qualified_id(citation.source, citation.chunk_id)] = ( + blobs + ) await chat_history.add_citations(citations, picture_bytes=picture_bytes) diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py index cb64245a..da42d634 100644 --- a/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py +++ b/haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py @@ -9,6 +9,7 @@ from textual.widgets import Collapsible, LoadingIndicator, Markdown, Static from textual.widgets.markdown import MarkdownStream from textual_image.widget import Image as TextualImage +from haiku.rag.store.models.chunk import qualified_id from haiku.rag.store.models.citation import Citation if TYPE_CHECKING: @@ -430,11 +431,12 @@ class ChatHistory(VerticalScroll): async def add_citations( self, citations: list[Citation], - picture_bytes: dict[str, list[bytes]] | None = None, + picture_bytes: dict[tuple[str | None, str | None], list[bytes]] | None = None, ) -> None: """Add citations inline after a response. - ``picture_bytes`` maps citation ``chunk_id`` → list of raw PNG bytes, + ``picture_bytes`` maps a citation's ``(source, chunk_id)`` → list of raw + PNG bytes, one per entry in the citation's ``picture_refs``. Pre-fetched by the caller (typically the chat app's post-response hook) so widget construction stays synchronous. @@ -445,7 +447,10 @@ class ChatHistory(VerticalScroll): picture_bytes = picture_bytes or {} for citation in citations: widget = CitationWidget( - citation, picture_bytes=picture_bytes.get(citation.chunk_id) + citation, + picture_bytes=picture_bytes.get( + qualified_id(citation.source, citation.chunk_id) + ), ) await self.mount(widget) self.scroll_end(animate=False) diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index ccf51c82..d9a02f56 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -40,8 +40,8 @@ def build_image_content_from_results( ) -> list[str | BinaryContent]: """Decode and validate picture bytes attached to search results, labelled. - Dedup keyed on ``(document_id, self_ref)`` so the same picture in - different chunks is sent once. Pictures that fail + Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in + different chunks is sent once, and a copy in another collection is its own. Pictures that fail ``PIL.Image.verify()`` are skipped — the model adapter renders one vision placeholder per ``BinaryContent``, so emitting one for an image the server can't decode leaves the processor with an @@ -58,12 +58,12 @@ def build_image_content_from_results( to the vision API. """ collected: list[tuple[str | None, str, BinaryContent]] = [] - seen: set[tuple[str | None, str]] = set() + seen: set[tuple[str | None, str | None, str]] = set() for result in results: if not result.image_data: continue for self_ref, b64 in result.image_data.items(): - key = (result.document_id, self_ref) + key = (result.source, result.document_id, self_ref) if key in seen: continue picture = decode_picture(base64.b64decode(b64), self_ref) diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index 03217592..710de420 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -137,6 +137,32 @@ def test_evidence_from_one_collection_does_not_name_it(): assert '[a] Source: "Title a" (test://a)' in lines +def test_a_picture_in_two_collections_is_retained_from_each(): + """A document copied into another collection keeps its id and picture refs.""" + found = discovered( + cited={"a": [2], "b": [2]}, + pictures={"a": ["#/pictures/0"], "b": ["#/pictures/0"]}, + ) + found = replace( + found, + citations={ + "a": replace_citation( + found.citations["a"], document_id="shared", source="papers" + ), + "b": replace_citation( + found.citations["b"], document_id="shared", source="wiki" + ), + }, + ) + + capsule = build_capsule([found]) + + assert [(picture.source, picture.self_ref) for picture in capsule.pictures] == [ + ("papers", "#/pictures/0"), + ("wiki", "#/pictures/0"), + ] + + def test_nothing_cited_produces_no_capsule(): capsule = build_capsule([discovered()]) diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index e2120c02..5408195c 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -732,6 +732,79 @@ class TestRenderingUnattributedPictures: covering.reader_for.assert_awaited_once_with(None) covering.get_picture_bytes.assert_not_awaited() + @pytest.mark.asyncio + async def test_one_chunk_id_in_two_collections_keeps_its_own_pictures( + self, temp_db_path: Path, monkeypatch + ): + """Chunk ids repeat between copies of a database, and both capabilities' + citations are gathered into one mapping.""" + from io import BytesIO + + from PIL import Image as PILImage + + from haiku.rag.chat.app import ANALYSIS_STATE_NAMESPACE, RAG_STATE_NAMESPACE + from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget + from haiku.rag.store.models.citation import Citation + + def png(color: str) -> bytes: + buffer = BytesIO() + PILImage.new("RGB", (4, 4), color).save(buffer, format="PNG") + return buffer.getvalue() + + pictures = {"alpha": png("red"), "beta": png("blue")} + + def reader(source): + owner = AsyncMock() + owner.get_picture_bytes = AsyncMock(return_value=pictures[source]) + return owner + + covering = _make_mock_client() + covering.covers_multiple = True + covering.source_names = ("alpha", "beta") + covering.reader_for = AsyncMock(side_effect=reader) + + def cited(source: str) -> dict: + return Citation( + document_id="d1", + chunk_id="c1", + source=source, + content="body", + document_uri=f"test://{source}", + picture_refs=["#/pictures/0"], + ).model_dump(mode="json") + + seen: list[tuple[str | None, list[bytes] | None]] = [] + build = CitationWidget.__init__ + + def capture(self, citation, picture_bytes=None, **kwargs): + seen.append((citation.source, picture_bytes)) + build(self, citation, picture_bytes, **kwargs) + + monkeypatch.setattr(CitationWidget, "__init__", capture) + + app, _ = _make_app(temp_db_path, covering) + with ( + patch("haiku.rag.chat.app.HaikuRAG") as stub, + _covering_returns(stub, covering), + ): + async with app.run_test() as pilot: + app._state[RAG_STATE_NAMESPACE] = { + "citations": ["c1"], + "citation_index": {"c1": cited("alpha")}, + } + app._state[ANALYSIS_STATE_NAMESPACE] = { + "citations": ["c1"], + "citation_index": {"c1": cited("beta")}, + } + + await app._show_citations_and_programs(app.query_one(ChatHistory)) + await pilot.pause() + + assert seen == [ + ("alpha", [pictures["alpha"]]), + ("beta", [pictures["beta"]]), + ] + class TestKeepingSelectionsReachable: """A selection applies whether or not the page shows it, and a checkbox is diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 38bb3fe4..4746c861 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -304,6 +304,37 @@ class TestBuildImageContentFromResults: images = [item for item in content if isinstance(item, BinaryContent)] assert len(images) == 1 + def test_the_same_picture_in_two_collections_is_attached_from_each(self): + """A document copied into another collection keeps its id and picture refs.""" + from pydantic_ai.messages import BinaryContent + + from haiku.rag.tools.search import build_image_content_from_results + + shared = {"#/pictures/0": _png_b64()} + results = [ + SearchResult( + content="a", + score=0.9, + chunk_id="c1", + document_id="doc-1", + source="papers", + image_data=shared, + ), + SearchResult( + content="b", + score=0.8, + chunk_id="c2", + document_id="doc-1", + source="wiki", + image_data=shared, + ), + ] + + content = build_image_content_from_results(results) + + images = [item for item in content if isinstance(item, BinaryContent)] + assert len(images) == 2 + def test_each_image_is_labelled_with_the_result_it_belongs_to(self): """Label every picture, not just the batch.