search: only attach picture bytes from pre-expansion chunks

haiku_rag_slim/haiku/rag/client/search.py

@@ -204,9 +204,10 @@ async def expand_context(
         expanded_results.extend(expanded)
 
     expanded_results.sort(key=lambda r: r.score, reverse=True)
-    # expand_with_items rebuilds SearchResult objects, so attach picture bytes
-    # to the fresh set — picture self_refs may have grown via section expansion.
-    await _populate_image_data(client, expanded_results)
+    # image_data and picture_captions are preserved through expansion by
+    # expand_with_items — we deliberately do not re-attach bytes for refs
+    # introduced by section expansion, so the multimodal payload stays
+    # bounded by what was originally retrieved.
     return expanded_results
 
 

haiku_rag_slim/haiku/rag/context.py

@@ -251,6 +251,20 @@ async def expand_with_items(
             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
+        # only pictures from the originally retrieved chunks get attached.
+        # Pictures swept in by section expansion are referenced in `refs`
+        # for cross-referencing but their bytes are not re-fetched —
+        # 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]
 
         # Expansion should never return less content than the original chunk.
@@ -272,6 +286,8 @@ async def expand_with_items(
                 page_numbers=sorted(pages) or first.page_numbers,
                 headings=all_headings or None,
                 labels=sorted(labels) or first.labels,
+                image_data=merged_image_data or None,
+                picture_captions=merged_captions,
             )
         )
 

tests/test_context.py

@@ -386,3 +386,108 @@ class TestExpandWithItems:
             # Expansion produces "Steps\n\nClick\n\n+\n\nAdd a New Service" = 38 chars
             # which is less than the chunk's 46 chars — fallback preserves the chunk
             assert expanded[0].content == result.content
+
+
+@pytest.mark.asyncio
+class TestExpandWithItemsPictureBytes:
+    """Picture bytes only ride along for refs present in the pre-expansion
+    chunk. Pictures swept in by section expansion are still referenced in
+    ``doc_item_refs`` for cross-referencing but their image_data is not
+    re-fetched — keeps the multimodal payload bounded.
+    """
+
+    async def _populate(self, rag):
+        items = [
+            DocumentItem(
+                document_id="doc-1",
+                position=0,
+                self_ref="#/texts/0",
+                label="section_header",
+                text="Section 1",
+            ),
+            DocumentItem(
+                document_id="doc-1",
+                position=1,
+                self_ref="#/texts/1",
+                label="text",
+                text="Paragraph in section 1.",
+            ),
+            DocumentItem(
+                document_id="doc-1",
+                position=2,
+                self_ref="#/pictures/0",
+                label="picture",
+                text="Figure 1 caption.",
+            ),
+            DocumentItem(
+                document_id="doc-1",
+                position=3,
+                self_ref="#/texts/2",
+                label="text",
+                text="Another paragraph after the figure.",
+            ),
+        ]
+        await rag.document_item_repository.create_items("doc-1", items)
+
+    async def test_pre_expansion_picture_bytes_preserved(self, temp_db_path):
+        """A picture chunk's image_data survives expansion."""
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as rag:
+            await self._populate(rag)
+            result = SearchResult(
+                content="Figure 1 caption.",
+                score=0.9,
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/0"],
+                image_data={"#/pictures/0": "BASE64BYTES"},
+                picture_captions={"#/pictures/0": "Figure 1 caption."},
+            )
+            expanded = await expand_with_items(
+                rag.document_item_repository, "doc-1", [result], 5000
+            )
+            assert len(expanded) == 1
+            assert expanded[0].image_data == {"#/pictures/0": "BASE64BYTES"}
+            assert expanded[0].picture_captions == {"#/pictures/0": "Figure 1 caption."}
+
+    async def test_merged_results_union_image_data(self, temp_db_path):
+        """When two results' ranges merge, their pre-expansion image_data
+        is unioned onto the merged output."""
+        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"#/pictures/{i}" if i in (1, 3) else f"#/texts/{i}",
+                    label="picture" if i in (1, 3) else "text",
+                    text=f"Item {i}.",
+                )
+                for i in range(5)
+            ]
+            await rag.document_item_repository.create_items("doc-1", items)
+
+            r1 = SearchResult(
+                content="Item 1.",
+                score=0.9,
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/1"],
+                image_data={"#/pictures/1": "A"},
+            )
+            r2 = SearchResult(
+                content="Item 3.",
+                score=0.85,
+                document_id="doc-1",
+                doc_item_refs=["#/pictures/3"],
+                image_data={"#/pictures/3": "B"},
+            )
+            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].image_data == {
+                "#/pictures/1": "A",
+                "#/pictures/3": "B",
+            }

tests/test_picture_in_context.py

@@ -203,10 +203,12 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
 
 
 @pytest.mark.asyncio
-async def test_expand_context_repopulates_image_data(temp_db_path):
-    """expand_context rebuilds SearchResult objects via expand_with_items, so
-    it must re-attach picture bytes — otherwise vision flows downstream see
-    empty image_data after expansion."""
+async def test_expand_context_does_not_attach_expansion_added_pictures(temp_db_path):
+    """expand_context preserves picture bytes from the pre-expansion result and
+    does NOT re-fetch bytes for picture self_refs swept in by section
+    expansion. The expansion-added picture ref still rides along in
+    doc_item_refs for cross-referencing, but image_data stays empty so the
+    multimodal payload is bounded by what search originally returned."""
     async with HaikuRAG(temp_db_path, create=True) as rag:
         await rag.document_item_repository.create_items(
             "doc-1",
@@ -247,7 +249,7 @@ async def test_expand_context_repopulates_image_data(temp_db_path):
         assert len(expanded) == 1
         out = expanded[0]
         assert "#/pictures/0" in out.doc_item_refs
-        assert out.image_data == {"#/pictures/0": PICTURE_B64}
+        assert out.image_data is None
 
 
 @dataclass