From fa319596ccff3eb009d56275834da055b1a22c84 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Aug 2026 12:02:16 +0300 Subject: [PATCH] Name the collection on a retrieved image, not only in the text Search results and capsule entries name the collection they came from; the images attached beside them carried only the chunk id and reference. Two collections can return the same picture of the same document, so the two labels were identical and the model could place neither. The decision is the one already made for the text: `covers_multiple` at the generic search tool, and the flag `search_corpus` computed for the capability tools, which it now returns. --- .../haiku/rag/capabilities/_base.py | 8 +- .../haiku/rag/capabilities/_tools.py | 6 +- .../haiku/rag/capabilities/compaction.py | 12 ++- haiku_rag_slim/haiku/rag/tools/search.py | 14 +-- tests/capabilities/test_capabilities.py | 53 +++++++++++ tests/capabilities/test_evidence_capsule.py | 44 ++++++++++ tests/multi_db/test_capabilities.py | 6 +- tests/tools/test_search.py | 87 ++++++++++++++++--- 8 files changed, 204 insertions(+), 26 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index b92f4b6d..ef55c3f3 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -505,7 +505,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): "the results you already have." ) async with self.rag_lock: - formatted, results = await search_corpus( + formatted, results, include_collection = await search_corpus( await self._ensure_rag(), query, limit=limit, @@ -517,7 +517,11 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): # narrower return must not drop what the wider one already showed it. merge_results(state.searches.setdefault(query, []), results) self._note_evidence() - if self.vision and (parts := build_image_content_from_results(results)): + if self.vision and ( + parts := build_image_content_from_results( + results, include_collection=include_collection + ) + ): return ToolReturn(return_value=formatted, content=parts) return formatted diff --git a/haiku_rag_slim/haiku/rag/capabilities/_tools.py b/haiku_rag_slim/haiku/rag/capabilities/_tools.py index 3d636988..a2964fd8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_tools.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_tools.py @@ -19,8 +19,8 @@ async def search_corpus( limit: int | None = None, document_filter: str | None = None, sources: list[str] | None = None, -) -> tuple[str, list[SearchResult]]: - """Search and context-expand results for a capability tool.""" +) -> tuple[str, list[SearchResult], bool]: + """Search and context-expand results, and whether they name their collection.""" results = await rag.search( query, limit=limit, filter=document_filter, sources=sources ) @@ -35,7 +35,7 @@ async def search_corpus( ) for index, result in enumerate(results) ) - return formatted or "No results found.", list(results) + return formatted or "No results found.", list(results), include_collection def merge_results( diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 4d77a378..84f083ba 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -49,10 +49,12 @@ def group_label(position: int) -> str: return f"[Cited evidence group {position}]" -def picture_label(chunk_id: str, self_ref: str) -> str: +def picture_label(chunk_id: str, self_ref: str, collection: str | None = None) -> str: + named = f"Collection: {collection}. " if collection else "" return ( f"Page image retrieved from the knowledge base for cited evidence " - f"[{chunk_id}] ({self_ref}). Not provided by the user. {RETRIEVED_IMAGE_TAG}" + f"[{chunk_id}] ({self_ref}). {named}" + f"Not provided by the user. {RETRIEVED_IMAGE_TAG}" ) @@ -191,7 +193,11 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: chunk_id=entry.chunk_id, document_id=entry.citation.document_id, self_ref=self_ref, - label=picture_label(entry.chunk_id, self_ref), + label=picture_label( + entry.chunk_id, + self_ref, + entry.citation.source if include_collection else None, + ), source=entry.citation.source, ) ) diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index d9a02f56..1ad4978e 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -37,6 +37,7 @@ def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None: def build_image_content_from_results( results: list[SearchResult], + include_collection: bool = False, ) -> list[str | BinaryContent]: """Decode and validate picture bytes attached to search results, labelled. @@ -57,7 +58,7 @@ def build_image_content_from_results( ``BinaryContent.identifier`` cannot do — it does not survive serialization to the vision API. """ - collected: list[tuple[str | None, str, BinaryContent]] = [] + collected: list[tuple[str | None, str | None, str, BinaryContent]] = [] seen: set[tuple[str | None, str | None, str]] = set() for result in results: if not result.image_data: @@ -69,15 +70,16 @@ def build_image_content_from_results( picture = decode_picture(base64.b64decode(b64), self_ref) if picture is None: continue - collected.append((result.chunk_id, self_ref, picture)) + collected.append((result.source, result.chunk_id, self_ref, picture)) seen.add(key) content: list[str | BinaryContent] = [] total = len(collected) - for position, (chunk_id, self_ref, picture) in enumerate(collected, 1): + for position, (source, chunk_id, self_ref, picture) in enumerate(collected, 1): + collection = f"Collection: {source}. " if include_collection and source else "" content.append( f"Page image {position} of {total}, retrieved from the knowledge base " - f"for search result [{chunk_id}] ({self_ref}). " + f"for search result [{chunk_id}] ({self_ref}). {collection}" f"Not provided by the user. {RETRIEVED_IMAGE_TAG}" ) content.append(picture) @@ -172,7 +174,9 @@ def create_search_toolset( if not config.qa.model.vision: return text - image_content = build_image_content_from_results(results_list) + image_content = build_image_content_from_results( + results_list, include_collection=include_collection + ) if image_content: return ToolReturn(return_value=text, content=image_content) return text diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 4b500494..f06856ab 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -425,6 +425,59 @@ def _stub_client(*batches: list[SearchResult]) -> AsyncMock: return client +def _picture_result(source: str) -> SearchResult: + import base64 + from io import BytesIO + + from PIL import Image as PILImage + + buffer = BytesIO() + PILImage.new("RGB", (4, 4), "red").save(buffer, format="PNG") + return SearchResult( + content="body", + score=0.9, + source=source, + chunk_id="c1", + document_id="d1", + image_data={"#/pictures/0": base64.b64encode(buffer.getvalue()).decode()}, + ) + + +async def _labels_of_search(temp_db_path, *sources: str) -> list[str]: + """The labels the search tool attaches to its images, for a client covering + `sources`.""" + from pydantic_ai.messages import ToolReturn + + capability = create_rag(db_path=temp_db_path, config=AppConfig(), vision=True) + capability.state = RAGState() + client = _stub_client([_picture_result(sources[0])]) + client.source_names = sources + + with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)): + returned = await capability._search("cats", None) + + assert isinstance(returned, ToolReturn) + assert returned.content is not None + return [item for item in returned.content if isinstance(item, str)] + + +@pytest.mark.asyncio +async def test_a_search_spanning_collections_names_them_on_its_images(temp_db_path): + """Images travel beside the results and are labelled the same way.""" + labels = await _labels_of_search(temp_db_path, "alpha", "beta") + + assert "Collection: alpha." in labels[0] + + +@pytest.mark.asyncio +async def test_a_search_over_one_collection_does_not_name_it_on_its_images( + temp_db_path, +): + labels = await _labels_of_search(temp_db_path, "alpha") + + assert not [label for label in labels if "Collection" in label] + + @pytest.mark.asyncio async def test_a_fruitless_search_says_so(temp_db_path): """A blank tool return reads as a broken tool, not as an empty corpus.""" diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index 710de420..6e3923ab 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -161,6 +161,50 @@ def test_a_picture_in_two_collections_is_retained_from_each(): ("papers", "#/pictures/0"), ("wiki", "#/pictures/0"), ] + # Nothing else tells the two apart once they are attached. + assert "Collection: papers." in capsule.pictures[0].label + assert "Collection: wiki." in capsule.pictures[1].label + + +def test_one_chunk_id_cited_from_two_collections_labels_each_picture(): + """Both capabilities can cite the same id from different collections, where + the reference and the document are the same too.""" + found = [ + replace( + discovered( + capability=capability, + cited={"c1": [2]}, + pictures={"c1": ["#/pictures/0"]}, + ), + citations={ + "c1": replace_citation( + citation("c1", pictures=["#/pictures/0"]), + document_id="shared", + source=source, + ) + }, + ) + for capability, source in (("analysis", "papers"), ("rag", "wiki")) + ] + + capsule = build_capsule(found) + + labels = [picture.label for picture in capsule.pictures] + assert all("[c1] (#/pictures/0)" in label for label in labels) + assert "Collection: papers." in labels[0] + assert "Collection: wiki." in labels[1] + + +def test_a_picture_from_one_collection_is_not_labelled_with_it(): + found = discovered(cited={"c1": [2]}, pictures={"c1": ["#/pictures/0"]}) + found = replace( + found, + citations={"c1": replace_citation(found.citations["c1"], source="papers")}, + ) + + [picture] = build_capsule([found]).pictures + + assert "Collection" not in picture.label def test_nothing_cited_produces_no_capsule(): diff --git a/tests/multi_db/test_capabilities.py b/tests/multi_db/test_capabilities.py index fd530fe9..a3a50422 100644 --- a/tests/multi_db/test_capabilities.py +++ b/tests/multi_db/test_capabilities.py @@ -306,11 +306,13 @@ class TestWhenTheModelIsToldTheCollection: async with HaikuRAG(config=config) as rag: monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha)) - spanning, _ = await search_corpus(rag, "cats") - narrowed, _ = await search_corpus(rag, "cats", sources=["alpha"]) + spanning, _, spans = await search_corpus(rag, "cats") + narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"]) assert "Collection: alpha" in spanning assert "Collection" not in narrowed + # Images travel beside the results and are labelled the same way. + assert (spans, narrows) == (True, False) class TestActionableFailures: diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 4746c861..5c0df56d 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -42,7 +42,7 @@ class TestNamingTheCollection: what the search spans.""" @staticmethod - def _client(covers_multiple: bool, source: str | None): + def _client(covers_multiple: bool, source: str | None, *, pictures: bool = False): from unittest.mock import AsyncMock results = [ @@ -53,6 +53,7 @@ class TestNamingTheCollection: chunk_id="c1", document_id="d1", document_title="Report", + image_data={"#/pictures/0": _png_b64()} if pictures else None, ) ] return SimpleNamespace( @@ -79,6 +80,42 @@ class TestNamingTheCollection: assert "Collection" not in text + @staticmethod + def _seeing_config(search_config): + config = search_config.model_copy(deep=True) + config.qa.model.vision = True + return config + + @pytest.mark.asyncio + async def test_a_client_covering_a_set_names_each_image(self, search_config): + """Images travel beside the results and are labelled the same way.""" + from pydantic_ai.messages import ToolReturn + + toolset = create_search_toolset(self._seeing_config(search_config)) + client = self._client(covers_multiple=True, source="alpha", pictures=True) + + returned = await toolset.tools["search"].function(make_ctx(client), "cats") + + assert isinstance(returned, ToolReturn) + labels = [item for item in returned.content if isinstance(item, str)] + assert "Collection: alpha." in labels[0] + + @pytest.mark.asyncio + async def test_one_named_collection_is_not_named_on_an_image(self, search_config): + from pydantic_ai.messages import ToolReturn + + toolset = create_search_toolset(self._seeing_config(search_config)) + client = self._client(covers_multiple=False, source="alpha", pictures=True) + + returned = await toolset.tools["search"].function(make_ctx(client), "cats") + + assert isinstance(returned, ToolReturn) + assert not [ + item + for item in returned.content + if isinstance(item, str) and "Collection" in item + ] + @pytest.mark.vcr() class TestSearchToolExecution: @@ -304,14 +341,11 @@ class TestBuildImageContentFromResults: images = [item for item in content if isinstance(item, BinaryContent)] assert len(images) == 1 - def test_the_same_picture_in_two_collections_is_attached_from_each(self): - """A document copied into another collection keeps its id and picture refs.""" - from pydantic_ai.messages import BinaryContent - - from haiku.rag.tools.search import build_image_content_from_results - + @staticmethod + def _one_picture_in_two_collections(): + """A document copied into another collection keeps its ids and refs.""" shared = {"#/pictures/0": _png_b64()} - results = [ + return [ SearchResult( content="a", score=0.9, @@ -321,20 +355,51 @@ class TestBuildImageContentFromResults: image_data=shared, ), SearchResult( - content="b", + content="a", score=0.8, - chunk_id="c2", + chunk_id="c1", document_id="doc-1", source="wiki", image_data=shared, ), ] - content = build_image_content_from_results(results) + def test_the_same_picture_in_two_collections_is_attached_from_each(self): + from pydantic_ai.messages import BinaryContent + + from haiku.rag.tools.search import build_image_content_from_results + + content = build_image_content_from_results( + self._one_picture_in_two_collections() + ) images = [item for item in content if isinstance(item, BinaryContent)] assert len(images) == 2 + def test_each_image_is_labelled_with_the_collection_it_came_from(self): + """Nothing else tells the two apart: same chunk id, same document, same + reference.""" + from haiku.rag.tools.search import build_image_content_from_results + + content = build_image_content_from_results( + self._one_picture_in_two_collections(), include_collection=True + ) + + labels = [item for item in content if isinstance(item, str)] + assert "Collection: papers." in labels[0] + assert "Collection: wiki." in labels[1] + + def test_an_unasked_for_collection_is_not_named_on_an_image(self): + from haiku.rag.tools.search import build_image_content_from_results + + content = build_image_content_from_results( + self._one_picture_in_two_collections() + ) + + assert not [ + item for item in content if isinstance(item, str) and "Collection" in item + ] + def test_each_image_is_labelled_with_the_result_it_belongs_to(self): """Label every picture, not just the batch.