From 46f7ab8d97c78233d8c76b65270e33b6bb8d5b3e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 8 Aug 2026 17:51:06 +0300 Subject: [PATCH] Name the search result each page image belongs to ToolReturn.content reaches the model as a user-role message and the pictures arrive bare, so nothing connects a figure to the chunk it came from: BinaryContent.identifier does not survive serialization to the vision API, and the captions in the result text correlate only by position. Precede each picture with its position, source chunk id and self_ref. build_binary_parts_from_results becomes build_image_content_from_results and returns the labels interleaved with the pictures, so both attachment sites emit them the same way. This does not stop a model narrating retrieved pictures as user-supplied. Measured on gemma4-26b with a single note ahead of the batch, and again with per-image labels: it quotes the label and still says the user provided them. The message role wins over its text. --- CHANGELOG.md | 1 + .../haiku/rag/capabilities/_base.py | 4 +- .../haiku/rag/store/models/chunk.py | 2 +- haiku_rag_slim/haiku/rag/tools/search.py | 46 +++++++---- tests/test_picture_in_context.py | 12 +-- tests/tools/test_search.py | 82 +++++++++++++++---- 6 files changed, 109 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f75350e7..5bda89f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ - `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`. - `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`. - `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized. +- Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures. - Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. - Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. - `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 41994ed8..5be8e73e 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -27,7 +27,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation, resolve_citations -from haiku.rag.tools.search import build_binary_parts_from_results +from haiku.rag.tools.search import build_image_content_from_results CITATION_GRACE_REQUESTS = 2 """Requests calling this capability's tools that its cite tool outlives the rest by. @@ -372,7 +372,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): ) state = cast(Any, self.state) state.searches[query] = results - if self.vision and (parts := build_binary_parts_from_results(results)): + if self.vision and (parts := build_image_content_from_results(results)): return ToolReturn(return_value=formatted, content=parts) return formatted diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index b2936e15..95666a39 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -211,7 +211,7 @@ class SearchResult(BaseModel): parts.append(f"Type: {primary_label}") # Surface picture captions when present. Order matches the binary - # attachments emitted by build_binary_parts_from_results, so the model + # attachments emitted by build_image_content_from_results, so the model # can correlate caption ↔ attached image by position (BinaryContent # identifiers don't survive serialization to the OpenAI vision API). if self.picture_captions: diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 4e8bf8e5..787210b6 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -11,10 +11,10 @@ from haiku.rag.store.models import SearchResult from haiku.rag.tools.context import RAGDeps -def build_binary_parts_from_results( +def build_image_content_from_results( results: list[SearchResult], -) -> list[BinaryContent]: - """Decode and validate picture bytes attached to search 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 @@ -22,8 +22,18 @@ def build_binary_parts_from_results( vision placeholder per ``BinaryContent``, so emitting one for an image the server can't decode leaves the processor with an off-by-one count. + + Every picture is preceded by a line naming the result it belongs to. + ``ToolReturn.content`` reaches the model as a user-role message, so + retrieved pictures are otherwise indistinguishable from ones the user + attached, and models narrate them as part of the question: unlabelled, + gemma4-26b answered about a figure from an unrelated document, and with a + single note ahead of the batch it still called them "images in the prompt". + The label also names the chunk to cite for a figure, which + ``BinaryContent.identifier`` cannot do — it does not survive serialization + to the vision API. """ - parts: list[BinaryContent] = [] + collected: list[tuple[str | None, str, bytes]] = [] seen: set[tuple[str | None, str]] = set() for result in results: if not result.image_data: @@ -38,15 +48,21 @@ def build_binary_parts_from_results( img.verify() except Exception: continue - parts.append( - BinaryContent( - data=data, - media_type="image/png", - identifier=self_ref, - ) - ) + collected.append((result.chunk_id, self_ref, data)) seen.add(key) - return parts + + content: list[str | BinaryContent] = [] + total = len(collected) + for position, (chunk_id, self_ref, data) in enumerate(collected, 1): + content.append( + f"Page image {position} of {total}, retrieved from the knowledge " + f"base for search result [{chunk_id}] ({self_ref}). " + "Not provided by the user." + ) + content.append( + BinaryContent(data=data, media_type="image/png", identifier=self_ref) + ) + return content def create_search_toolset( @@ -134,9 +150,9 @@ def create_search_toolset( if not config.qa.model.vision: return text - binary_parts = build_binary_parts_from_results(results_list) - if binary_parts: - return ToolReturn(return_value=text, content=binary_parts) + image_content = build_image_content_from_results(results_list) + if image_content: + return ToolReturn(return_value=text, content=image_content) return text toolset: FunctionToolset[RAGDeps] = FunctionToolset() diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 55f2f368..a9b32a47 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -402,8 +402,9 @@ async def test_search_tool_returns_multimodal_when_picture_present(): assert isinstance(result.return_value, str) assert "Type: picture" in result.return_value or "rank 1" in result.return_value assert result.content is not None - assert len(result.content) == 1 - part = result.content[0] + images = [c for c in result.content if isinstance(c, BinaryContent)] + assert len(images) == 1 + part = images[0] assert isinstance(part, BinaryContent) assert part.media_type == "image/png" assert part.identifier == "#/pictures/0" @@ -457,11 +458,12 @@ async def test_search_tool_attaches_same_self_ref_from_different_documents(): assert isinstance(result, ToolReturn) assert result.content is not None - assert len(result.content) == 2, ( + images = [c for c in result.content if isinstance(c, BinaryContent)] + assert len(images) == 2, ( "Both documents' figures must reach the model — dedup keyed on " "self_ref alone would drop doc-B's bytes." ) - payloads = {part.data for part in result.content} # type: ignore[attr-defined] + payloads = {part.data for part in images} assert PICTURE_BYTES in payloads assert other_bytes in payloads @@ -905,7 +907,7 @@ async def test_search_tool_drops_invalid_image_bytes(): assert isinstance(result, ToolReturn) assert result.content is not None - identifiers = {p.identifier for p in result.content} # type: ignore[attr-defined] + identifiers = {p.identifier for p in result.content if isinstance(p, BinaryContent)} assert identifiers == {"#/pictures/0"}, ( "Only the decodable PNG should reach the model — the corrupt " "ref must be dropped so we don't emit a placeholder for an " diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 60bf3404..8f09ad04 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -210,30 +210,35 @@ def search_config(): return Config -class TestBuildBinaryPartsFromResults: - """Picture bytes are attached once per (document, self_ref) pair.""" +def _png_b64(): + import base64 + from io import BytesIO + + from PIL import Image as PILImage + + buf = BytesIO() + PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode() + + +class TestBuildImageContentFromResults: + """Picture bytes are attached once per (document, self_ref) pair, and labelled.""" def test_results_without_image_data_contribute_nothing(self): - from haiku.rag.tools.search import build_binary_parts_from_results + from haiku.rag.tools.search import build_image_content_from_results results = [ SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None) ] - assert build_binary_parts_from_results(results) == [] + assert build_image_content_from_results(results) == [] def test_duplicate_document_and_ref_is_attached_once(self): - import base64 - from io import BytesIO + from pydantic_ai.messages import BinaryContent - from PIL import Image as PILImage + from haiku.rag.tools.search import build_image_content_from_results - from haiku.rag.tools.search import build_binary_parts_from_results - - buf = BytesIO() - PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG") - png = base64.b64encode(buf.getvalue()).decode() - shared = {"#/pictures/0": png} + shared = {"#/pictures/0": _png_b64()} results = [ SearchResult( content="a", @@ -251,6 +256,53 @@ class TestBuildBinaryPartsFromResults: ), ] - parts = build_binary_parts_from_results(results) + content = build_image_content_from_results(results) - assert len(parts) == 1 + images = [item for item in content if isinstance(item, BinaryContent)] + assert len(images) == 1 + + def test_each_image_is_labelled_with_the_result_it_belongs_to(self): + """Label every picture, not just the batch. + + ``ToolReturn.content`` reaches the model as a user-role message, and one + leading note does not override that: with a single note on the wire, + gemma4-26b still reasoned "the user also provided images in the prompt". + A label adjacent to each picture also names the chunk to cite for it, + which ``BinaryContent.identifier`` cannot do — it does not survive + serialization to the vision API. + """ + from pydantic_ai.messages import BinaryContent + + from haiku.rag.tools.search import build_image_content_from_results + + results = [ + SearchResult( + content="a", + score=0.9, + chunk_id="c1", + document_id="doc-1", + image_data={"#/pictures/0": _png_b64()}, + ), + SearchResult( + content="b", + score=0.8, + chunk_id="c2", + document_id="doc-2", + image_data={"#/pictures/3": _png_b64()}, + ), + ] + + content = build_image_content_from_results(results) + + # label, image, label, image — each picture preceded by its own line. + assert [type(item) is str for item in content] == [True, False, True, False] + assert isinstance(content[1], BinaryContent) + assert isinstance(content[3], BinaryContent) + + first, second = content[0], content[2] + assert isinstance(first, str) and isinstance(second, str) + assert "c1" in first and "#/pictures/0" in first + assert "c2" in second and "#/pictures/3" in second + assert "1 of 2" in first and "2 of 2" in second + for label in (first, second): + assert "not provided by the user" in label.lower()