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