From e3314e36fdc48392ea8ddb31658ce1f3d85bf5cf Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 19 May 2026 14:49:16 +0300 Subject: [PATCH] search: only attach picture bytes from pre-expansion chunks --- haiku_rag_slim/haiku/rag/client/search.py | 7 +- haiku_rag_slim/haiku/rag/context.py | 16 ++++ tests/test_context.py | 105 ++++++++++++++++++++++ tests/test_picture_in_context.py | 12 +-- 4 files changed, 132 insertions(+), 8 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 495887ff..8cb2ab50 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index f32b14e5..38eb531c 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/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, ) ) diff --git a/tests/test_context.py b/tests/test_context.py index 0ad6665a..2c807a53 100644 --- a/tests/test_context.py +++ b/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", + } diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 9c7e2cf1..7c959b77 100644 --- a/tests/test_picture_in_context.py +++ 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