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