From ddac328d05b00f229022c6eaea9c4fdc098af179 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 2 Sep 2026 11:35:17 +0300 Subject: [PATCH 1/3] Price qa.max_searches in search units Searches a model emits in one response share a budget unit, up to FREE_SIBLINGS_PER_ROUND (3) per unit; sequential searches pay one unit each, as before. Grouping keys on RunContext.run_step, which pydantic-ai increments once per model request. A budget-rejected round fails all its remaining siblings, and tracking resets per run. Glimmer opens most questions with a burst of ~3 rephrasings in a single response (95.8% of its three-search ORB cases are one-response bursts), spending 3 of 5 searches before reading anything. Pass rate at 3 calls equals 1 call, so a burst is priced as one probe. --- CHANGELOG.md | 6 + .../haiku/rag/capabilities/_base.py | 31 +++- .../haiku/rag/capabilities/analysis.py | 2 +- haiku_rag_slim/haiku/rag/capabilities/rag.py | 2 +- tests/capabilities/test_capabilities.py | 18 +- tests/capabilities/test_citation_policy.py | 2 +- tests/capabilities/test_evidence_wire.py | 4 +- tests/capabilities/test_search_units.py | 156 ++++++++++++++++++ tests/multi_db/test_capabilities.py | 8 +- tests/multi_db/test_citations.py | 4 +- tests/test_picture_in_context.py | 2 +- 11 files changed, 213 insertions(+), 22 deletions(-) create mode 100644 tests/capabilities/test_search_units.py diff --git a/CHANGELOG.md b/CHANGELOG.md index dadbc52a..f5ada525 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## [Unreleased] +### Changed + +- `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. + ## [0.81.0] - 2026-09-01 ### Added diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index d9ca411b..94b1d592 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -61,6 +61,15 @@ Calibration knob. Two unrelated UUID4s reach about 0.5, while dropping or duplicating a character or a whole group stays above 0.75, so the gap is wide. """ +FREE_SIBLINGS_PER_ROUND = 3 +"""Searches one budget unit covers when emitted in the same model response. + +Calibration knob, sized to the measured modal burst. ``qa.max_searches`` +counts units, so a model rephrasing its query a few times in one response +spends one unit, while every search of a sequential searcher is a unit of its +own. +""" + def _ambiguous_retry(error: AmbiguousCitationError) -> ModelRetry: """The only way out of an id that names a chunk in two databases. @@ -177,6 +186,10 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) search_count: int = field(default=0, repr=False) + search_step: int = field(default=0, repr=False) + """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) request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False) @@ -226,6 +239,9 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): rag_lock=asyncio.Lock(), resource_lock=asyncio.Lock(), search_count=0, + search_step=0, + step_searches=0, + step_rejected=False, request_count=0, grace_requests_used=0, epoch=0, @@ -493,10 +509,19 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): retrieved_now=retrieved, ) - async def _search(self, query: str, limit: int | None) -> str | ToolReturn: + async def _search( + self, query: str, limit: int | None, run_step: int + ) -> str | ToolReturn: assert self.state is not None - self.search_count += 1 - if self.search_count > self._max_searches: + if run_step != self.search_step: + self.search_step = run_step + self.step_searches = 0 + self.step_rejected = False + self.step_searches += 1 + if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0: + self.search_count += 1 + if self.step_rejected or self.search_count > self._max_searches: + self.step_rejected = True raise ToolFailed( "Search limit reached. Answer the question using " "the results you already have." diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index fb856ccc..fb88a081 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -172,7 +172,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ctx: RunContext[Any], query: str, limit: int | None = None ) -> str | ToolReturn: """Search the knowledge base for evidence to analyze.""" - return await self._with_state(self._search(query, limit)) + return await self._with_state(self._search(query, limit, ctx.run_step)) async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any: """Execute Python against the sandboxed document filesystem.""" diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index e4dfb138..6da47837 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -81,7 +81,7 @@ class RAGCapability(RAGCapabilityBase[RAGState]): ctx: RunContext[Any], query: str, limit: int | None = None ) -> str | ToolReturn: """Search the knowledge base using hybrid vector and full-text search.""" - return await self._with_state(self._search(query, limit)) + return await self._with_state(self._search(query, limit, ctx.run_step)) async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any: """Register exact search-result chunk IDs as citations for the answer.""" diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 424fc36c..a3202542 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -412,7 +412,7 @@ async def test_a_spent_search_budget_fails_the_tool(temp_db_path): capability.state = RAGState() with pytest.raises(ToolFailed, match="Search limit reached"): - await capability._search("anything", None) + await capability._search("anything", None, 1) def _stub_client(*batches: list[SearchResult]) -> AsyncMock: @@ -452,7 +452,7 @@ async def _labels_of_search(temp_db_path, *sources: str) -> list[str]: client.source_names = sources with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)): - returned = await capability._search("cats", None) + returned = await capability._search("cats", None, 1) assert isinstance(returned, ToolReturn) assert returned.content is not None @@ -483,7 +483,9 @@ async def test_a_fruitless_search_says_so(temp_db_path): capability.state = RAGState() capability.borrowed_rag = _stub_client([]) - assert await capability._search("nothing about this", None) == "No results found." + assert ( + await capability._search("nothing about this", None, 1) == "No results found." + ) @pytest.mark.asyncio @@ -500,8 +502,8 @@ async def test_a_narrower_repeat_keeps_what_the_wider_search_returned(temp_db_pa [SearchResult(content="first", score=1.0, chunk_id="chunk-1")], ) - await capability._search("Figure 3-1", 20) - await capability._search("Figure 3-1", None) + await capability._search("Figure 3-1", 20, 1) + await capability._search("Figure 3-1", None, 2) stored = capability.state.searches["Figure 3-1"] assert [result.chunk_id for result in stored] == [ @@ -525,8 +527,8 @@ async def test_two_databases_holding_one_chunk_id_both_survive(temp_db_path): ], ) - await capability._search("cats", 20) - await capability._search("cats", None) + await capability._search("cats", 20, 1) + await capability._search("cats", None, 2) stored = capability.state.searches["cats"] assert [(r.source, r.chunk_id) for r in stored] == [ @@ -1105,7 +1107,7 @@ def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord: return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) -async def _stub_search(self, query: str, _limit: int | None) -> str: +async def _stub_search(self, query: str, _limit: int | None, _run_step: int) -> str: """Record a result the way the real search does, so citing resolves.""" cast(Any, self.state).searches[query] = [ SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") diff --git a/tests/capabilities/test_citation_policy.py b/tests/capabilities/test_citation_policy.py index bcc3fde9..2a16c9fb 100644 --- a/tests/capabilities/test_citation_policy.py +++ b/tests/capabilities/test_citation_policy.py @@ -29,7 +29,7 @@ class Deps: state: dict[str, Any] = field(default_factory=dict) -async def stub_search(self, query: str, _limit: int | None) -> str: +async def stub_search(self, query: str, _limit: int | None, _run_step: int) -> str: cast(Any, self.state).searches[query] = [ SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") ] diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index bc80ea43..263a0e35 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -437,7 +437,9 @@ REAL_PNG = base64.b64decode( ) -async def _search_with_a_picture(self, query: str, _limit: int | None) -> str: +async def _search_with_a_picture( + self, query: str, _limit: int | None, _run_step: int +) -> str: """Record a result carrying a page image, the way a real search does.""" cast(Any, self.state).searches[query] = [ SearchResult( diff --git a/tests/capabilities/test_search_units.py b/tests/capabilities/test_search_units.py new file mode 100644 index 00000000..f55f7676 --- /dev/null +++ b/tests/capabilities/test_search_units.py @@ -0,0 +1,156 @@ +from dataclasses import dataclass, field +from typing import Any + +import pytest +from pydantic_ai import Agent +from pydantic_ai.messages import ( + ModelResponse, + TextPart, + ToolCallPart, + ToolReturnPart, +) +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.run import AgentRunResult + +from haiku.rag.capabilities.rag import create_capability as create_rag +from haiku.rag.config.models import AppConfig + + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + +def burst_model(bursts: list[list[str]]) -> FunctionModel: + """Emit one `rag_search` call per query in each burst, then answer.""" + responses = 0 + + def model_function(_messages, _info) -> ModelResponse: + nonlocal responses + responses += 1 + if responses <= len(bursts): + return ModelResponse( + parts=[ + ToolCallPart("rag_search", {"query": query}) + for query in bursts[responses - 1] + ] + ) + return ModelResponse(parts=[TextPart("done")]) + + return FunctionModel(model_function) + + +def burst_agent( + bursts: list[list[str]], db_path, max_searches: int +) -> Agent[Deps, str]: + config = AppConfig() + config.qa.max_searches = max_searches + return Agent( + burst_model(bursts), + deps_type=Deps, + capabilities=[create_rag(db_path=db_path, config=config, defer_loading=False)], + ) + + +def search_returns(result: AgentRunResult[Any]) -> list[ToolReturnPart]: + return [ + part + for message in result.all_messages() + for part in message.parts + if isinstance(part, ToolReturnPart) and part.tool_name == "rag_search" + ] + + +def outcomes(result: AgentRunResult[Any]) -> list[str]: + return [ + "failed" if part.outcome == "failed" else "ok" + for part in search_returns(result) + ] + + +@pytest.mark.asyncio +async def test_a_burst_in_one_response_consumes_one_unit(rag_db): + """Three searches emitted together cost one unit and run in emission order.""" + agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 1) + + result = await agent.run("question", deps=Deps()) + + assert outcomes(result) == ["ok", "ok", "ok"] + calls = [ + part + for message in result.all_messages() + for part in message.parts + if isinstance(part, ToolCallPart) and part.tool_name == "rag_search" + ] + assert [part.tool_call_id for part in search_returns(result)] == [ + part.tool_call_id for part in calls + ] + + +@pytest.mark.asyncio +async def test_sequential_searches_pay_one_unit_each(rag_db): + agent = burst_agent([["ai"], ["machine learning"]], rag_db, 1) + + result = await agent.run("question", deps=Deps()) + + assert outcomes(result) == ["ok", "failed"] + assert "Search limit reached" in str(search_returns(result)[1].content) + + +@pytest.mark.asyncio +async def test_max_searches_zero_fails_every_sibling(rag_db): + agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 0) + + result = await agent.run("question", deps=Deps()) + + assert outcomes(result) == ["failed", "failed", "failed"] + + +@pytest.mark.asyncio +async def test_a_rejected_round_fails_all_its_siblings(rag_db): + agent = burst_agent([["ai"], ["ml", "deep learning", "supervised"]], rag_db, 1) + + result = await agent.run("question", deps=Deps()) + + assert outcomes(result) == ["ok", "failed", "failed", "failed"] + + +@pytest.mark.asyncio +async def test_a_sibling_past_the_allowance_pays_its_own_unit(rag_db): + burst = [["ai", "machine learning", "deep learning", "supervised learning"]] + + within = await burst_agent(burst, rag_db, 2).run("question", deps=Deps()) + over = await burst_agent(burst, rag_db, 1).run("question", deps=Deps()) + + assert outcomes(within) == ["ok", "ok", "ok", "ok"] + assert outcomes(over) == ["ok", "ok", "ok", "failed"] + + +@pytest.mark.asyncio +async def test_unit_tracking_resets_between_runs(rag_db): + """A second run's opening burst prices like a first run's.""" + + def model_function(messages, _info) -> ModelResponse: + if any(isinstance(part, ToolReturnPart) for part in messages[-1].parts): + return ModelResponse(parts=[TextPart("done")]) + return ModelResponse( + parts=[ + ToolCallPart("rag_search", {"query": query}) + for query in ["ai", "machine learning", "deep learning"] + ] + ) + + config = AppConfig() + config.qa.max_searches = 1 + agent = Agent( + FunctionModel(model_function), + deps_type=Deps, + capabilities=[create_rag(db_path=rag_db, config=config, defer_loading=False)], + ) + deps = Deps() + + first = await agent.run("question", deps=deps) + second = await agent.run("another", deps=deps, message_history=first.all_messages()) + + assert outcomes(first) == ["ok", "ok", "ok"] + assert outcomes(second)[-3:] == ["ok", "ok", "ok"] diff --git a/tests/multi_db/test_capabilities.py b/tests/multi_db/test_capabilities.py index 77a55c91..b67f3fd8 100644 --- a/tests/multi_db/test_capabilities.py +++ b/tests/multi_db/test_capabilities.py @@ -30,7 +30,7 @@ class TestAskAcrossDatabases: capability = create_capability(config=config, rag=rag, defer_loading=False) capability.state = RAGState(sources=["alpha"]) - formatted = await capability._search("cats", limit=10) + formatted = await capability._search("cats", 10, 1) assert isinstance(formatted, str) assert "alpha" in formatted @@ -47,7 +47,7 @@ class TestAskAcrossDatabases: capability = create_capability(config=config, rag=rag, defer_loading=False) capability.state = RAGState() - formatted = await capability._search("cats", limit=10) + formatted = await capability._search("cats", 10, 1) assert isinstance(formatted, str) assert "alpha document" in formatted @@ -72,7 +72,7 @@ class TestStandaloneCapabilities: assert capability.scope.names == ("alpha", "beta") run = await capability.for_run(make_context(Deps())) try: - formatted = await run._search("cats", limit=10) + formatted = await run._search("cats", 10, 1) finally: await run._close() @@ -135,7 +135,7 @@ class TestAnalyzeAcrossDatabases: capability = create_analysis(config=config, rag=rag, defer_loading=False) capability.state = AnalysisState(sources=["alpha"]) - formatted = await capability._search("cats", limit=10) + formatted = await capability._search("cats", 10, 1) sandbox = await capability._ensure_sandbox() await capability._close() diff --git a/tests/multi_db/test_citations.py b/tests/multi_db/test_citations.py index 73e238b9..62f5d223 100644 --- a/tests/multi_db/test_citations.py +++ b/tests/multi_db/test_citations.py @@ -342,7 +342,7 @@ class TestCiteFallback: ) run = await capability.for_run(make_context(deps)) # The search returns the cats chunk, never the aardvark one. - await run._search("cats", limit=10) + await run._search("cats", 10, 1) await run._cite([aardvark.id]) @@ -374,7 +374,7 @@ class TestCiteFallback: state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")} ) run = await capability.for_run(make_context(deps)) - await run._search("cats", limit=10) + await run._search("cats", 10, 1) with pytest.raises(ModelRetry, match="None of the supplied chunk_ids"): await run._cite([outside.id]) diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 3f464797..1f1cbd2d 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -1022,7 +1022,7 @@ async def test_rag_capability_attaches_images_for_vision_model(temp_db_path): capability.state = RAGState() capability.rag = fake_client - result = await capability._search("anything", None) + result = await capability._search("anything", None, 1) assert isinstance(result, ToolReturn) assert result.content is not None From d9f489dcc843ff0211e926114a4ee0ab54dedfc8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 2 Sep 2026 12:05:20 +0300 Subject: [PATCH 2/3] 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. --- CHANGELOG.md | 3 + .../haiku/rag/capabilities/_base.py | 33 +++- .../haiku/rag/capabilities/_tools.py | 69 ++++++- haiku_rag_slim/haiku/rag/tools/search.py | 34 +++- tests/capabilities/test_search_units.py | 173 ++++++++++++++++++ tests/multi_db/test_capabilities.py | 6 +- tests/tools/test_search.py | 12 +- 7 files changed, 300 insertions(+), 30 deletions(-) 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] From 31fc8aba295c7ff93ae705c8c97b1549e47dc65c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 2 Sep 2026 12:20:33 +0300 Subject: [PATCH 3/3] Guard dedup against compaction and document search units A picture deduplicated within a burst must still reach the model once in its own question and, when cited, again via the capsule re-fetch after compaction; uncited it is dropped. Two wire tests pin both paths. The harness builds the rag capability with defer_loading=False: a deferred evidence capability the model has not loaded presents an empty record, so compaction computes boundary 0 and rewrites nothing. qa.md describes max_searches in units. --- docs/configuration/qa.md | 4 +- tests/capabilities/test_evidence_wire.py | 83 ++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 2 deletions(-) diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index 9716c4d1..05e4c5a4 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -30,12 +30,12 @@ qa: enable_thinking: true temperature: 0.3 # Default: 0.3 vision: false # Set true for vision-capable models - max_searches: 5 # Maximum search tool calls per question + max_searches: 5 # Maximum search units per question ``` - **model**: LLM configuration (see [Providers](providers.md#model-settings)) - **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The capability's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix. -- **max_searches**: Maximum number of search tool calls a capability can make per question (default: 5). Shared by the RAG and analysis capabilities. +- **max_searches**: Maximum number of search units a capability can spend per question (default: 5). Up to three searches emitted in the same model response share one unit, so a model that rephrases its query in one response spends one unit. A search in a later response starts a new unit, as does each further group of three within one response. Shared by the RAG and analysis capabilities. Searches in one response also deduplicate their returns: evidence a sibling search already showed collapses to a reference line, and each picture attaches once per response. !!! note "Thinking on vLLM" `enable_thinking` only applies to models with a pydantic-ai reasoning profile (o-series, gpt-5, gpt-oss). For other vLLM-served models such as Qwen3 or the Gemma family, the field is a silent no-op — set the chat template switch via [`extra_body`](providers.md#raw-provider-pass-through) instead. diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 263a0e35..c189c2f6 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -517,6 +517,89 @@ async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label( assert texts_of(wire[-1]) == [] +def _burst_result() -> SearchResult: + return SearchResult( + content="evidence", + score=1.0, + chunk_id="chunk-1", + document_id="doc-1", + source="main", + doc_item_refs=["#/pictures/0"], + image_data={"#/pictures/0": base64.b64encode(REAL_PNG).decode()}, + ) + + +async def _fanout_question_then_another(temp_db_path, cite: bool) -> list[list[Any]]: + """Question 1 fans out over one picture chunk; question 2 follows compacted.""" + rag = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False, vision=True + ) + client = AsyncMock() + client.search.side_effect = [[_burst_result()], [_burst_result()]] + client.expand_context.side_effect = lambda results: results + client.source_names = ["main"] + rag.borrowed_rag = client + citing = ( + [[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-3")]] + if cite + else [] + ) + calls = iter( + [ + [ + ToolCallPart("rag_search", {"query": "figure"}, "call-1"), + ToolCallPart("rag_search", {"query": "the figure"}, "call-2"), + ], + *citing, + [TextPart("first answer")], + [TextPart("second answer")], + ] + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=next(calls)) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag, create_compaction()], + ) + deps = Deps() + + with patch.object( + RAGCapability, "get_picture_bytes", AsyncMock(return_value=REAL_PNG) + ): + first = await agent.run("what does the figure show?", deps=deps) + await agent.run( + "and what else?", deps=deps, message_history=first.all_messages() + ) + return wire + + +@pytest.mark.asyncio +async def test_a_burst_deduplicated_picture_survives_compaction_when_cited( + temp_db_path, +): + """Dedup attaches the picture once in its own question; the capsule re-fetches + it for the next. Neither pass may leave the model without it.""" + wire = await _fanout_question_then_another(temp_db_path, cite=True) + + assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG] + assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG] + + +@pytest.mark.asyncio +async def test_a_burst_deduplicated_picture_is_dropped_by_compaction_uncited( + temp_db_path, +): + wire = await _fanout_question_then_another(temp_db_path, cite=False) + + assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG] + assert images_of(wire[-1]) == [] + + @pytest.mark.asyncio async def test_the_capsule_is_built_once_per_request_and_again_for_the_next( temp_db_path,