Deduplicate search results within one model response

Sibling searches emitted in one response overlap heavily (40.6% of
returned chunk slots on Glimmer ORB fan-out cases). A result whose
rendered evidence a sibling already showed keeps its rank slot but
collapses to a reference line, and a picture attaches once per response
keyed on (source, document_id, self_ref). Equivalence is the
format_for_agent rendering at neutral rank/total plus picture keys,
bucketed under the qualified chunk id, so another database's copy or a
different expansion of the same anchor formats in full.

Search state now commits only after formatting and image construction
succeed: a raising image build no longer leaves results citable that
the model never saw, notes evidence for them, or suppresses a later
sibling.
This commit is contained in:
Yiorgis Gozadinos 2026-09-02 12:05:20 +03:00
parent ddac328d05
commit d9f489dcc8
No known key found for this signature in database
7 changed files with 300 additions and 30 deletions

View file

@ -7,6 +7,9 @@
- `qa.max_searches` counts search units: searches a model emits in one - `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 response share a unit, up to 3 per unit; sequential searches pay one unit
each. 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 ## [0.81.0] - 2026-09-01

View file

@ -14,6 +14,7 @@ from pydantic_ai import (
) )
from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ( from pydantic_ai.messages import (
BinaryContent,
InstructionPart, InstructionPart,
ModelMessage, ModelMessage,
ModelRequest, ModelRequest,
@ -29,6 +30,7 @@ from pydantic_ai.toolsets import AgentToolset
from haiku.rag.capabilities._tools import ( from haiku.rag.capabilities._tools import (
CodeExecutionEntry, CodeExecutionEntry,
EvidenceKey,
merge_results, merge_results,
search_corpus, search_corpus,
) )
@ -43,7 +45,7 @@ from haiku.rag.store.models.citation import (
ambiguous_citation, ambiguous_citation,
resolve_citations, 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 CITATION_GRACE_REQUESTS = 2
"""Requests calling this capability's tools that its cite tool outlives the rest by. """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.""" """The run_step whose searches are being priced and deduplicated."""
step_searches: int = field(default=0, repr=False) step_searches: int = field(default=0, repr=False)
step_rejected: bool = field(default=False, 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) request_count: int = field(default=0, repr=False)
grace_requests_used: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False)
epoch: 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, search_step=0,
step_searches=0, step_searches=0,
step_rejected=False, step_rejected=False,
step_shown=set(),
step_pictures=set(),
request_count=0, request_count=0,
grace_requests_used=0, grace_requests_used=0,
epoch=0, epoch=0,
@ -517,6 +523,8 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
self.search_step = run_step self.search_step = run_step
self.step_searches = 0 self.step_searches = 0
self.step_rejected = False self.step_rejected = False
self.step_shown = set()
self.step_pictures = set()
self.step_searches += 1 self.step_searches += 1
if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0: if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0:
self.search_count += 1 self.search_count += 1
@ -527,23 +535,34 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
"the results you already have." "the results you already have."
) )
async with self.rag_lock: async with self.rag_lock:
formatted, results, include_collection = await search_corpus( formatted, results, rendered, include_collection = await search_corpus(
await self._ensure_rag(), await self._ensure_rag(),
query, query,
limit=limit, limit=limit,
document_filter=self.state.document_filter, document_filter=self.state.document_filter,
sources=self.state.sources, 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 state = self.state
# A model can search the same query twice with different limits, and the # 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. # narrower return must not drop what the wider one already showed it.
merge_results(state.searches.setdefault(query, []), results) merge_results(state.searches.setdefault(query, []), results)
self._note_evidence() self._note_evidence()
if self.vision and ( self.step_shown |= rendered
parts := build_image_content_from_results( self.step_pictures |= emitted
results, include_collection=include_collection if parts:
)
):
return ToolReturn(return_value=formatted, content=parts) return ToolReturn(return_value=formatted, content=parts)
return formatted return formatted

View file

@ -1,9 +1,11 @@
from collections.abc import Iterable from collections.abc import Iterable
from collections.abc import Set as AbstractSet
from pydantic import BaseModel from pydantic import BaseModel
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.chunk import SearchResult, qualified_id from haiku.rag.store.models.chunk import SearchResult, qualified_id
from haiku.rag.tools.search import picture_keys
class CodeExecutionEntry(BaseModel): class CodeExecutionEntry(BaseModel):
@ -13,14 +15,48 @@ class CodeExecutionEntry(BaseModel):
success: bool = True 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( async def search_corpus(
rag: HaikuRAG, rag: HaikuRAG,
query: str, query: str,
limit: int | None = None, limit: int | None = None,
document_filter: str | None = None, document_filter: str | None = None,
sources: list[str] | None = None, sources: list[str] | None = None,
) -> tuple[str, list[SearchResult], bool]: shown: AbstractSet[EvidenceKey] = frozenset(),
"""Search and context-expand results, and whether they name their collection.""" ) -> 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( results = await rag.search(
query, limit=limit, filter=document_filter, sources=sources 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. # two collections names them even when everything came back from one.
selected = rag.source_names if sources is None else sources selected = rag.source_names if sources is None else sources
include_collection = len(set(selected)) > 1 include_collection = len(set(selected)) > 1
formatted = "\n\n---\n\n".join( rendered: set[EvidenceKey] = set()
result.format_for_agent( parts: list[str] = []
rank=index + 1, total=len(results), include_collection=include_collection total = len(results)
) for index, result in enumerate(results):
for index, result in enumerate(results) key = evidence_key(result, include_collection)
) if key in shown or key in rendered:
return formatted or "No results found.", list(results), include_collection 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( def merge_results(
@ -56,6 +104,9 @@ def merge_results(
__all__ = [ __all__ = [
"CodeExecutionEntry", "CodeExecutionEntry",
"EvidenceKey",
"evidence_key",
"evidence_signature",
"merge_results", "merge_results",
"search_corpus", "search_corpus",
] ]

View file

@ -1,5 +1,6 @@
import base64 import base64
from collections.abc import Callable from collections.abc import Callable
from collections.abc import Set as AbstractSet
from io import BytesIO from io import BytesIO
from PIL import Image 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: def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
"""Wrap picture bytes for the wire, or return nothing if they will not decode. """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( def build_image_content_from_results(
results: list[SearchResult], results: list[SearchResult],
include_collection: bool = False, 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. """Decode and validate picture bytes attached to search results, labelled.
Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in Returns the labelled content and the ``PictureKey`` of every picture it
different chunks is sent once, and a copy in another collection is its own. Pictures that fail 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 ``PIL.Image.verify()`` are skipped the model adapter renders one
vision placeholder per ``BinaryContent``, so emitting one for an vision placeholder per ``BinaryContent``, so emitting one for an
image the server can't decode leaves the processor with 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. to the vision API.
""" """
collected: list[tuple[str | None, str | None, str, BinaryContent]] = [] 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: for result in results:
if not result.image_data: if not result.image_data:
continue continue
@ -72,6 +93,7 @@ def build_image_content_from_results(
continue continue
collected.append((result.source, result.chunk_id, self_ref, picture)) collected.append((result.source, result.chunk_id, self_ref, picture))
seen.add(key) seen.add(key)
emitted.add(key)
content: list[str | BinaryContent] = [] content: list[str | BinaryContent] = []
total = len(collected) total = len(collected)
@ -83,7 +105,7 @@ def build_image_content_from_results(
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}" f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
) )
content.append(picture) content.append(picture)
return content return content, emitted
def create_search_toolset( def create_search_toolset(
@ -174,7 +196,7 @@ def create_search_toolset(
if not config.qa.model.vision: if not config.qa.model.vision:
return text return text
image_content = build_image_content_from_results( image_content, _ = build_image_content_from_results(
results_list, include_collection=include_collection results_list, include_collection=include_collection
) )
if image_content: if image_content:

View file

@ -1,9 +1,14 @@
import base64
from dataclasses import dataclass, field from dataclasses import dataclass, field
from io import BytesIO
from typing import Any from typing import Any
from unittest.mock import AsyncMock
import pytest import pytest
from PIL import Image as PILImage
from pydantic_ai import Agent from pydantic_ai import Agent
from pydantic_ai.messages import ( from pydantic_ai.messages import (
BinaryContent,
ModelResponse, ModelResponse,
TextPart, TextPart,
ToolCallPart, ToolCallPart,
@ -12,8 +17,11 @@ from pydantic_ai.messages import (
from pydantic_ai.models.function import FunctionModel from pydantic_ai.models.function import FunctionModel
from pydantic_ai.run import AgentRunResult 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.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
@dataclass @dataclass
@ -154,3 +162,168 @@ async def test_unit_tracking_resets_between_runs(rag_db):
assert outcomes(first) == ["ok", "ok", "ok"] assert outcomes(first) == ["ok", "ok", "ok"]
assert outcomes(second)[-3:] == ["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

View file

@ -305,8 +305,10 @@ class TestWhenTheModelIsToldTheCollection:
async with HaikuRAG(config=config) as rag: async with HaikuRAG(config=config) as rag:
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha)) monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
spanning, _, spans = await search_corpus(rag, "cats") spanning, _, _, spans = await search_corpus(rag, "cats")
narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"]) narrowed, _, _, narrows = await search_corpus(
rag, "cats", sources=["alpha"]
)
assert "Collection: alpha" in spanning assert "Collection: alpha" in spanning
assert "Collection" not in narrowed assert "Collection" not in narrowed

View file

@ -311,7 +311,7 @@ class TestBuildImageContentFromResults:
SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None) 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): def test_duplicate_document_and_ref_is_attached_once(self):
from pydantic_ai.messages import BinaryContent 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)] images = [item for item in content if isinstance(item, BinaryContent)]
assert len(images) == 1 assert len(images) == 1
@ -369,7 +369,7 @@ class TestBuildImageContentFromResults:
from haiku.rag.tools.search import build_image_content_from_results 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() self._one_picture_in_two_collections()
) )
@ -381,7 +381,7 @@ class TestBuildImageContentFromResults:
reference.""" reference."""
from haiku.rag.tools.search import build_image_content_from_results 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 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): def test_an_unasked_for_collection_is_not_named_on_an_image(self):
from haiku.rag.tools.search import build_image_content_from_results 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() 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. # 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 [type(item) is str for item in content] == [True, False, True, False]