From c96894a0f0c15a29012f4d0a2f327b1bafc4b0e4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 5 May 2026 15:56:10 +0300 Subject: [PATCH] fix cross-document picture dedup in agent and skill search tools --- haiku_rag_slim/haiku/rag/skills/_tools.py | 7 +-- haiku_rag_slim/haiku/rag/tools/search.py | 7 +-- tests/test_picture_in_context.py | 56 +++++++++++++++++++++++ tests/test_skill_tools.py | 44 ++++++++++++++++++ 4 files changed, 108 insertions(+), 6 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index d5613e80..c5116c8c 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/haiku_rag_slim/haiku/rag/skills/_tools.py @@ -203,12 +203,13 @@ def create_skill_tools( return formatted binary_parts: list[BinaryContent] = [] - seen: set[str] = set() + seen: set[tuple[str | None, str]] = set() for result in results: if not result.image_data: continue for self_ref, b64 in result.image_data.items(): - if self_ref in seen: + key = (result.document_id, self_ref) + if key in seen: continue binary_parts.append( BinaryContent( @@ -217,7 +218,7 @@ def create_skill_tools( identifier=self_ref, ) ) - seen.add(self_ref) + seen.add(key) if binary_parts: return ToolReturn(return_value=formatted, content=binary_parts) diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 237ef327..5715227c 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -94,12 +94,13 @@ def create_search_toolset( return text binary_parts: list[BinaryContent] = [] - seen: set[str] = set() + seen: set[tuple[str | None, str]] = set() for result in results_list: if not result.image_data: continue for self_ref, b64 in result.image_data.items(): - if self_ref in seen: + key = (result.document_id, self_ref) + if key in seen: continue binary_parts.append( BinaryContent( @@ -108,7 +109,7 @@ def create_search_toolset( identifier=self_ref, ) ) - seen.add(self_ref) + seen.add(key) if binary_parts: return ToolReturn(return_value=text, content=binary_parts) diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index d52daf6d..bc3675eb 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -290,6 +290,62 @@ async def test_search_tool_returns_multimodal_when_picture_present(): assert part.data == PICTURE_BYTES +@pytest.mark.asyncio +async def test_search_tool_attaches_same_self_ref_from_different_documents(): + """Two different documents both have ``#/pictures/0`` — the dedup must + key on ``(document_id, self_ref)`` so each document's figure reaches + the model. Keying on ``self_ref`` alone silently drops the second + document's bytes, leaving the model with text only for that result.""" + other_bytes = b"\x89PNG\r\n\x1a\nother-doc-bytes" + other_b64 = base64.b64encode(other_bytes).decode("ascii") + + doc_a = SearchResult( + content="Figure from doc A", + score=1.0, + chunk_id="chunk-a", + document_id="doc-A", + doc_item_refs=["#/pictures/0"], + labels=["picture"], + image_data={"#/pictures/0": PICTURE_B64}, + ) + doc_b = SearchResult( + content="Figure from doc B (same self_ref, different bytes)", + score=0.9, + chunk_id="chunk-b", + document_id="doc-B", + doc_item_refs=["#/pictures/0"], + labels=["picture"], + image_data={"#/pictures/0": other_b64}, + ) + + fake_client = AsyncMock() + fake_client.search = AsyncMock(return_value=[doc_a, doc_b]) + fake_client.expand_context = AsyncMock(return_value=[doc_a, doc_b]) + + config = AppConfig() + config.qa.model.vision = True + toolset = create_search_toolset(config, expand_context=False) + func = toolset.tools["search"].function + + ctx = RunContext( + deps=_Deps(client=fake_client), # type: ignore[arg-type] + model=TestModel(), + usage=RunUsage(), + run_id="run-1", + ) + result = await func(ctx, "anything") + + assert isinstance(result, ToolReturn) + assert result.content is not None + assert len(result.content) == 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] + assert PICTURE_BYTES in payloads + assert other_bytes in payloads + + # Synthetic picture chunks at ingest diff --git a/tests/test_skill_tools.py b/tests/test_skill_tools.py index 49a1ad16..b6403e1c 100644 --- a/tests/test_skill_tools.py +++ b/tests/test_skill_tools.py @@ -185,3 +185,47 @@ async def test_skill_search_dedups_picture_bytes_by_self_ref(): assert result.content is not None assert len(result.content) == 1 assert result.content[0].identifier == "#/pictures/0" # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_skill_search_keeps_same_self_ref_from_different_documents(): + """``#/pictures/0`` in document A and ``#/pictures/0`` in document B + are different figures. Dedup must key on ``(document_id, self_ref)``; + keying on ``self_ref`` alone would drop document B's bytes.""" + other_bytes = b"\x89PNG\r\n\x1a\nother-doc-bytes" + other_b64 = base64.b64encode(other_bytes).decode("ascii") + + doc_a = SearchResult( + content="Figure in doc A", + score=1.0, + chunk_id="chunk-a", + document_id="doc-A", + doc_item_refs=["#/pictures/0"], + labels=["picture"], + image_data={"#/pictures/0": PICTURE_B64}, + ) + doc_b = SearchResult( + content="Figure in doc B", + score=0.9, + chunk_id="chunk-b", + document_id="doc-B", + doc_item_refs=["#/pictures/0"], + labels=["picture"], + image_data={"#/pictures/0": other_b64}, + ) + + config = AppConfig() + config.qa.model.vision = True + + search = _build_search_tool(config) + rag = _fake_rag([doc_a, doc_b]) + ctx = _make_ctx(rag, RAGState()) + + result = await search(ctx, "figure") + + assert isinstance(result, ToolReturn) + assert result.content is not None + assert len(result.content) == 2 + payloads = {part.data for part in result.content} # type: ignore[attr-defined] + assert PICTURE_BYTES in payloads + assert other_bytes in payloads