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.