diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index 7fd8e500..65331838 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -7,6 +7,7 @@ from pydantic_ai import Agent from pydantic_ai.messages import ( ModelMessage, ModelResponse, + RetryPromptPart, ToolCallPart, ToolReturnPart, ) @@ -57,6 +58,16 @@ def _count_tool_traffic( execution budget and any error in model-written Python both surface as ``ToolFailed``, distinguishable in the history only by message text — so ``n_failed_tools`` covers both without claiming to tell them apart. + + ``n_failed_tools`` counts ``RetryPromptPart`` as well as failed returns. + ``_cite`` rejects a call with ``ModelRetry`` rather than ``ToolFailed``, and + only the latter sets ``outcome="failed"``, so counting returns alone would + leave a run whose every cite attempt was rejected reporting zero failures. + + ``requests`` is the run's model-request count, which equals the capability's + own ``request_count`` budget only while the capability is loaded for the + whole run. A deferred capability skips hooks until it loads, so a router + turn before that is counted here and not against its limit. """ search_tool = f"{namespace}_search" search_calls = 0 @@ -73,7 +84,9 @@ def _count_tool_traffic( ) continue for part in message.parts: - if isinstance(part, ToolReturnPart) and part.outcome == "failed": + if isinstance(part, RetryPromptPart): + failed_tools += 1 + elif isinstance(part, ToolReturnPart) and part.outcome == "failed": failed_tools += 1 if part.tool_name == search_tool: rejected_searches += 1 diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index 8e3cf317..cd8c979e 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/evaluations/tests/test_capability_runner.py @@ -5,6 +5,7 @@ import pytest from pydantic_ai.messages import ( ModelRequest, ModelResponse, + RetryPromptPart, TextPart, ToolCallPart, ToolReturnPart, @@ -18,6 +19,31 @@ from haiku.rag.capabilities.rag import create_capability as create_rag from haiku.rag.config.models import AppConfig +def test_count_tool_traffic_sees_a_rejected_cite_call(): + """`_cite` rejects with ModelRetry, which is not a failed ToolReturnPart.""" + messages = [ + ModelRequest(parts=[UserPromptPart(content="q")]), + ModelResponse(parts=[ToolCallPart("analysis_cite", {"chunk_ids": []})]), + ModelRequest( + parts=[ + RetryPromptPart( + tool_name="analysis_cite", + content="No citations registered: chunk_ids was empty.", + tool_call_id="1", + ) + ] + ), + ModelResponse(parts=[TextPart("done")]), + ] + + _search_calls, rejected_searches, failed_tools, _requests = _count_tool_traffic( + messages, "analysis" + ) + + assert failed_tools == 1 + assert rejected_searches == 0 + + def test_count_tool_traffic_separates_search_rejections_from_code_errors(): """A crash in model-written Python must not read as budget exhaustion.""" messages = [ diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 9bd51bf9..603d224d 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -191,11 +191,20 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): "evidence already gathered." ) if spent := self._spent_tool_names(): + names = ", ".join(sorted(spent)) + if remaining := sorted(self._evidence_tool_names() - spent): + return ( + f"The {self.state_namespace} capability has spent its budget " + f"for {names}; further calls to them fail. Gather any further " + f"evidence with {', '.join(remaining)}, or call " + f"{self._cite_tool_name} with the chunk_ids you have and " + "answer." + ) return ( f"The {self.state_namespace} capability has spent its budget for " - f"{', '.join(sorted(spent))}; further calls to them fail. Answer " - f"from the evidence already gathered and call " - f"{self._cite_tool_name} with the chunk_ids supporting it." + f"{names}; further calls to them fail. Answer from the evidence " + f"already gathered and call {self._cite_tool_name} with the " + "chunk_ids supporting it." ) return None @@ -225,6 +234,14 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): if tool.capability_id != self.id or tool.name == self._cite_tool_name ] + def _evidence_tool_names(self) -> set[str]: + """Tools that can bring new evidence into the run. + + Drives the spent-budget notice: while one of these still has budget the + model must be pointed at it, not told to answer from what it has. + """ + return {f"{self.state_namespace}_search"} + def _spent_tool_names(self) -> set[str]: """This capability's tools whose own budget is exhausted.""" if self.search_count >= self._max_searches: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index cea8a878..95fa155e 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -85,6 +85,12 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): self.sandbox = None await super()._close() + def _evidence_tool_names(self) -> set[str]: + # Searching from inside the sandbox does not count against + # `qa.max_searches`, so code execution outlives a spent search budget + # as a way to reach new evidence. + return super()._evidence_tool_names() | {"analysis_execute_code"} + def _spent_tool_names(self) -> set[str]: spent = super()._spent_tool_names() if self.execute_count >= self.config.analysis.max_executions: diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index e062bd9c..715330a3 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -494,6 +494,46 @@ async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db): assert "spent its budget for rag_search" in seen_instructions[1] +@pytest.mark.asyncio +async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_path): + """Analysis must be sent to the sandbox, not told to answer, while it can. + + In-code `search()` bypasses `qa.max_searches`, and the instructions tell the + model to escalate to code when search results are insufficient. + """ + config = AppConfig() + config.qa.max_searches = 2 + capability = create_analysis(db_path=temp_db_path, config=config) + capability.search_count = 2 + + notice = capability._budget_notice() + + assert notice is not None + assert "analysis_search" in notice + assert "analysis_execute_code" in notice + assert "Answer from the evidence already gathered" not in notice + + # Once the code budget is gone too there is nowhere left to go. + capability.execute_count = config.analysis.max_executions + notice = capability._budget_notice() + assert notice is not None + assert "Answer from the evidence already gathered" in notice + + +@pytest.mark.asyncio +async def test_spent_search_notice_tells_rag_to_answer(temp_db_path): + """Search is the RAG capability's only evidence tool, so stopping is right.""" + config = AppConfig() + config.qa.max_searches = 2 + capability = create_rag(db_path=temp_db_path, config=config) + capability.search_count = 2 + + notice = capability._budget_notice() + + assert notice is not None + assert "Answer from the evidence already gathered" in notice + + @pytest.mark.asyncio async def test_spent_execution_budget_joins_the_notice(temp_db_path): config = AppConfig()