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_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
+    Addressed by owner, collection, document and reference, because a reference
-    ``#/pictures/0`` repeats across documents and capabilities. The label travels
+    such as ``#/pictures/0`` repeats across all three. The label travels with it
-    with it so it can never be emitted without its image.
+    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
+            # A reference such as `#/pictures/0` repeats across documents, and a
-            # in another document stays a different picture.
+            # document repeats across collections.
-            identity = (entry.capability, entry.citation.document_id, self_ref)
+            identity = (
+                entry.capability,
+                entry.citation.source,
+                entry.citation.document_id,
+                self_ref,
+            )
             if identity in seen:
                 continue
             seen.add(identity)

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)
 

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)

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
+    Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in
-    different chunks is sent once. Pictures that fail
+    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)

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()])
 

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

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.