diff --git a/CHANGELOG.md b/CHANGELOG.md index f5ada525..68b8b5c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ - `qa.max_searches` counts search units: searches a model emits in one response share a unit, up to 3 per unit; sequential searches pay one unit each. +- Searches in one model response deduplicate their results: evidence a sibling + search already showed collapses to a reference line, and a picture attaches + once per response. ## [0.81.0] - 2026-09-01 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 94b1d592..56c91c84 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -14,6 +14,7 @@ from pydantic_ai import ( ) from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.messages import ( + BinaryContent, InstructionPart, ModelMessage, ModelRequest, @@ -29,6 +30,7 @@ from pydantic_ai.toolsets import AgentToolset from haiku.rag.capabilities._tools import ( CodeExecutionEntry, + EvidenceKey, merge_results, search_corpus, ) @@ -43,7 +45,7 @@ from haiku.rag.store.models.citation import ( ambiguous_citation, resolve_citations, ) -from haiku.rag.tools.search import build_image_content_from_results +from haiku.rag.tools.search import PictureKey, build_image_content_from_results CITATION_GRACE_REQUESTS = 2 """Requests calling this capability's tools that its cite tool outlives the rest by. @@ -190,6 +192,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): """The run_step whose searches are being priced and deduplicated.""" step_searches: int = field(default=0, repr=False) step_rejected: bool = field(default=False, repr=False) + step_shown: set[EvidenceKey] = field(default_factory=set, repr=False) + step_pictures: set[PictureKey] = field(default_factory=set, repr=False) request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False) @@ -242,6 +246,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): search_step=0, step_searches=0, step_rejected=False, + step_shown=set(), + step_pictures=set(), request_count=0, grace_requests_used=0, epoch=0, @@ -517,6 +523,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): self.search_step = run_step self.step_searches = 0 self.step_rejected = False + self.step_shown = set() + self.step_pictures = set() self.step_searches += 1 if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0: self.search_count += 1 @@ -527,23 +535,34 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): "the results you already have." ) async with self.rag_lock: - formatted, results, include_collection = await search_corpus( + formatted, results, rendered, include_collection = await search_corpus( await self._ensure_rag(), query, limit=limit, document_filter=self.state.document_filter, sources=self.state.sources, + shown=self.step_shown, ) + parts: list[str | BinaryContent] = [] + emitted: set[PictureKey] = set() + if self.vision: + parts, emitted = build_image_content_from_results( + results, + include_collection=include_collection, + exclude=self.step_pictures, + ) + # Everything the search produced commits together, after formatting and + # image construction have both succeeded: a search that raises must not + # leave results citable, note evidence the model never received, or + # suppress a later sibling's results. state = self.state # A model can search the same query twice with different limits, and the # 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, include_collection=include_collection - ) - ): + self.step_shown |= rendered + self.step_pictures |= emitted + if parts: 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 a2964fd8..9450e27f 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_tools.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_tools.py @@ -1,9 +1,11 @@ from collections.abc import Iterable +from collections.abc import Set as AbstractSet from pydantic import BaseModel from haiku.rag.client import HaikuRAG from haiku.rag.store.models.chunk import SearchResult, qualified_id +from haiku.rag.tools.search import picture_keys class CodeExecutionEntry(BaseModel): @@ -13,14 +15,48 @@ class CodeExecutionEntry(BaseModel): success: bool = True +EvidenceKey = tuple[tuple[str | None, str | None], tuple[str, frozenset]] +"""What tells one rendered result from another: qualified id, then signature. + +The qualified id comes first because the rendered string alone would conflate +identical renderings of the same chunk id held by two databases. +""" + + +def evidence_signature(result: SearchResult, include_collection: bool) -> tuple: + """The rendered evidence a result shows the model, as an equivalence key. + + Rank and total are held at neutral values: they vary with a result's + position, and position (like score) must not tell two renderings apart. + """ + return ( + result.format_for_agent(rank=0, total=0, include_collection=include_collection), + picture_keys(result), + ) + + +def evidence_key(result: SearchResult, include_collection: bool) -> EvidenceKey: + return ( + qualified_id(result.source, result.chunk_id), + evidence_signature(result, include_collection), + ) + + async def search_corpus( rag: HaikuRAG, query: str, limit: int | None = None, document_filter: str | None = None, sources: list[str] | None = None, -) -> tuple[str, list[SearchResult], bool]: - """Search and context-expand results, and whether they name their collection.""" + shown: AbstractSet[EvidenceKey] = frozenset(), +) -> tuple[str, list[SearchResult], set[EvidenceKey], bool]: + """Search and context-expand results, eliding evidence already shown. + + Returns the formatted results, the full result list, the evidence keys the + formatting rendered in full, and whether results name their collection. A + result whose key is in ``shown`` keeps its slot but collapses to one line; + the result list is never filtered. + """ results = await rag.search( query, limit=limit, filter=document_filter, sources=sources ) @@ -29,13 +65,25 @@ async def search_corpus( # two collections names them even when everything came back from one. selected = rag.source_names if sources is None else sources include_collection = len(set(selected)) > 1 - formatted = "\n\n---\n\n".join( - result.format_for_agent( - rank=index + 1, total=len(results), include_collection=include_collection - ) - for index, result in enumerate(results) - ) - return formatted or "No results found.", list(results), include_collection + rendered: set[EvidenceKey] = set() + parts: list[str] = [] + total = len(results) + for index, result in enumerate(results): + key = evidence_key(result, include_collection) + if key in shown or key in rendered: + parts.append( + f"Also matched, shown above: [{result.chunk_id}] " + f"[rank {index + 1} of {total}]" + ) + else: + parts.append( + result.format_for_agent( + rank=index + 1, total=total, include_collection=include_collection + ) + ) + rendered.add(key) + formatted = "\n\n---\n\n".join(parts) + return formatted or "No results found.", list(results), rendered, include_collection def merge_results( @@ -56,6 +104,9 @@ def merge_results( __all__ = [ "CodeExecutionEntry", + "EvidenceKey", + "evidence_key", + "evidence_signature", "merge_results", "search_corpus", ] diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 1ad4978e..c779a19a 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -1,5 +1,6 @@ import base64 from collections.abc import Callable +from collections.abc import Set as AbstractSet from io import BytesIO from PIL import Image @@ -20,6 +21,22 @@ their own picture had it removed, along with their text. """ +PictureKey = tuple[str | None, str | None, str] +"""Identity of one attached picture: (source, document_id, self_ref). + +``self_ref`` alone collides across documents, and a copy of a document in +another collection carries its own pictures. +""" + + +def picture_keys(result: SearchResult) -> frozenset[PictureKey]: + """The identity of every picture this result carries.""" + return frozenset( + (result.source, result.document_id, self_ref) + for self_ref in (result.image_data or {}) + ) + + def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None: """Wrap picture bytes for the wire, or return nothing if they will not decode. @@ -38,11 +55,14 @@ 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]: + exclude: AbstractSet[PictureKey] = frozenset(), +) -> tuple[list[str | BinaryContent], set[PictureKey]]: """Decode and validate picture bytes attached to search results, labelled. - Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in - different chunks is sent once, and a copy in another collection is its own. Pictures that fail + Returns the labelled content and the ``PictureKey`` of every picture it + emitted. Dedup keyed on ``PictureKey`` so the same picture in + different chunks is sent once, and a copy in another collection is its + own; ``exclude`` seeds that dedup with pictures already sent. Pictures that fail ``PIL.Image.verify()`` are skipped — the model adapter renders one vision placeholder per ``BinaryContent``, so emitting one for an image the server can't decode leaves the processor with an @@ -59,7 +79,8 @@ def build_image_content_from_results( to the vision API. """ collected: list[tuple[str | None, str | None, str, BinaryContent]] = [] - seen: set[tuple[str | None, str | None, str]] = set() + seen: set[PictureKey] = set(exclude) + emitted: set[PictureKey] = set() for result in results: if not result.image_data: continue @@ -72,6 +93,7 @@ def build_image_content_from_results( continue collected.append((result.source, result.chunk_id, self_ref, picture)) seen.add(key) + emitted.add(key) content: list[str | BinaryContent] = [] total = len(collected) @@ -83,7 +105,7 @@ def build_image_content_from_results( f"Not provided by the user. {RETRIEVED_IMAGE_TAG}" ) content.append(picture) - return content + return content, emitted def create_search_toolset( @@ -174,7 +196,7 @@ def create_search_toolset( if not config.qa.model.vision: return text - image_content = build_image_content_from_results( + image_content, _ = build_image_content_from_results( results_list, include_collection=include_collection ) if image_content: diff --git a/tests/capabilities/test_search_units.py b/tests/capabilities/test_search_units.py index f55f7676..ed3db5e4 100644 --- a/tests/capabilities/test_search_units.py +++ b/tests/capabilities/test_search_units.py @@ -1,9 +1,14 @@ +import base64 from dataclasses import dataclass, field +from io import BytesIO from typing import Any +from unittest.mock import AsyncMock import pytest +from PIL import Image as PILImage from pydantic_ai import Agent from pydantic_ai.messages import ( + BinaryContent, ModelResponse, TextPart, ToolCallPart, @@ -12,8 +17,11 @@ from pydantic_ai.messages import ( from pydantic_ai.models.function import FunctionModel from pydantic_ai.run import AgentRunResult +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord +from haiku.rag.capabilities.rag import RAGState from haiku.rag.capabilities.rag import create_capability as create_rag from haiku.rag.config.models import AppConfig +from haiku.rag.store.models.chunk import SearchResult @dataclass @@ -154,3 +162,168 @@ async def test_unit_tracking_resets_between_runs(rag_db): assert outcomes(first) == ["ok", "ok", "ok"] assert outcomes(second)[-3:] == ["ok", "ok", "ok"] + + +def _png() -> str: + buffer = BytesIO() + PILImage.new("RGB", (4, 4), "red").save(buffer, format="PNG") + return base64.b64encode(buffer.getvalue()).decode() + + +def make_result(**overrides: Any) -> SearchResult: + fields: dict[str, Any] = { + "content": "body", + "score": 0.9, + "source": "main", + "chunk_id": "c1", + "document_id": "d1", + "image_data": {"#/pictures/0": _png()}, + } + fields.update(overrides) + return SearchResult(**fields) + + +def stub_client( + *batches: list[SearchResult], sources: list[str] | None = None +) -> AsyncMock: + client = AsyncMock() + client.search.side_effect = list(batches) + client.expand_context.side_effect = lambda results: results + client.source_names = sources or ["main"] + return client + + +def dedup_capability(client: AsyncMock, temp_db_path, *, vision: bool = True): + capability = create_rag(db_path=temp_db_path, config=AppConfig(), vision=vision) + capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) + capability.borrowed_rag = client + return capability + + +def images_of(returned: Any) -> list[BinaryContent]: + if isinstance(returned, str): + return [] + return [item for item in returned.content if isinstance(item, BinaryContent)] + + +def text_of(returned: Any) -> str: + return returned if isinstance(returned, str) else returned.return_value + + +@pytest.mark.asyncio +async def test_a_duplicate_sibling_is_elided_and_stays_citable(temp_db_path): + duplicate, novel = make_result(), make_result(chunk_id="c2", content="novel") + client = stub_client([make_result()], [duplicate, novel]) + capability = dedup_capability(client, temp_db_path) + + first = await capability._search("q", None, 1) + second = await capability._search("q rephrased", None, 1) + + assert len(images_of(first)) == 1 + assert images_of(second) == [] + text = text_of(second) + assert "Also matched, shown above: [c1] [rank 1 of 2]" in text + assert "body" not in text + assert "[rank 2 of 2]" in text and "novel" in text + assert [r.chunk_id for r in capability.state.searches["q rephrased"]] == [ + "c1", + "c2", + ] + assert await capability._cite(["c1"]) == "Registered 1 citation(s)." + + +@pytest.mark.asyncio +async def test_a_new_run_step_formats_shown_results_in_full(temp_db_path): + client = stub_client([make_result()], [make_result()]) + capability = dedup_capability(client, temp_db_path) + + await capability._search("q", None, 1) + second = await capability._search("q again", None, 2) + + assert "body" in text_of(second) + assert len(images_of(second)) == 1 + + +@pytest.mark.asyncio +async def test_same_chunk_id_from_another_collection_is_not_elided(temp_db_path): + client = stub_client( + [make_result(source="alpha")], + [make_result(source="beta")], + sources=["alpha", "beta"], + ) + capability = dedup_capability(client, temp_db_path) + + await capability._search("q", None, 1) + second = await capability._search("q rephrased", None, 1) + + assert "body" in text_of(second) + + +@pytest.mark.asyncio +async def test_same_anchor_with_new_evidence_formats_in_full(temp_db_path): + shared, extra = _png(), _png() + client = stub_client( + [make_result(content="c1 with c2", image_data={"#/pictures/1": shared})], + [ + make_result( + content="c1 with c3", + image_data={"#/pictures/1": shared, "#/pictures/3": extra}, + ) + ], + ) + capability = dedup_capability(client, temp_db_path) + + await capability._search("q", None, 1) + second = await capability._search("q rephrased", None, 1) + + assert "c1 with c3" in text_of(second) + assert len(images_of(second)) == 1 + labels = [item for item in second.content if isinstance(item, str)] + assert any("#/pictures/3" in label for label in labels) + + +@pytest.mark.parametrize( + ("overrides", "elided"), + [ + ({"score": 0.1}, True), + ({"content": "different"}, False), + ({"document_title": "Other"}, False), + ({"headings": ["Heading"]}, False), + ({"labels": ["table"]}, False), + ({"picture_captions": {"#/pictures/0": "A caption"}}, False), + ({"image_data": {"#/pictures/9": _png()}}, False), + ], +) +@pytest.mark.asyncio +async def test_equivalence_follows_the_rendered_evidence( + temp_db_path, overrides: dict[str, Any], elided: bool +): + """Any rendered field or picture identity defeats elision; score alone does not.""" + client = stub_client([make_result()], [make_result(**overrides)]) + capability = dedup_capability(client, temp_db_path) + + await capability._search("q", None, 1) + second = await capability._search("q rephrased", None, 1) + + assert ("Also matched, shown above" in text_of(second)) is elided + + +@pytest.mark.asyncio +async def test_a_failed_sibling_commits_nothing(temp_db_path): + client = stub_client( + [make_result(image_data={"#/pictures/0": "AAA"})], + [make_result()], + ) + capability = dedup_capability(client, temp_db_path) + evidence_before = capability.state.evidence.model_dump() + + with pytest.raises(Exception): + await capability._search("q", None, 1) + + assert capability.state.searches == {} + assert capability.state.evidence.model_dump() == evidence_before + + second = await capability._search("q rephrased", None, 1) + + assert "body" in text_of(second) + assert len(images_of(second)) == 1 diff --git a/tests/multi_db/test_capabilities.py b/tests/multi_db/test_capabilities.py index b67f3fd8..192af9b8 100644 --- a/tests/multi_db/test_capabilities.py +++ b/tests/multi_db/test_capabilities.py @@ -305,8 +305,10 @@ class TestWhenTheModelIsToldTheCollection: async with HaikuRAG(config=config) as rag: monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha)) - spanning, _, spans = await search_corpus(rag, "cats") - narrowed, _, narrows = 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 diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 5c0df56d..1723df8d 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -311,7 +311,7 @@ class TestBuildImageContentFromResults: SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None) ] - assert build_image_content_from_results(results) == [] + assert build_image_content_from_results(results) == ([], set()) def test_duplicate_document_and_ref_is_attached_once(self): from pydantic_ai.messages import BinaryContent @@ -336,7 +336,7 @@ class TestBuildImageContentFromResults: ), ] - content = build_image_content_from_results(results) + content, _ = build_image_content_from_results(results) images = [item for item in content if isinstance(item, BinaryContent)] assert len(images) == 1 @@ -369,7 +369,7 @@ class TestBuildImageContentFromResults: from haiku.rag.tools.search import build_image_content_from_results - content = build_image_content_from_results( + content, _ = build_image_content_from_results( self._one_picture_in_two_collections() ) @@ -381,7 +381,7 @@ class TestBuildImageContentFromResults: reference.""" from haiku.rag.tools.search import build_image_content_from_results - content = build_image_content_from_results( + content, _ = build_image_content_from_results( self._one_picture_in_two_collections(), include_collection=True ) @@ -392,7 +392,7 @@ class TestBuildImageContentFromResults: 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( + content, _ = build_image_content_from_results( self._one_picture_in_two_collections() ) @@ -431,7 +431,7 @@ class TestBuildImageContentFromResults: ), ] - content = build_image_content_from_results(results) + 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]