From 721acbcf3802a9d19200a4ff0dedc833895eff47 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 30 Jul 2026 15:50:06 +0300 Subject: [PATCH] Address review on PR #524 - _budget_notice no longer names the cite tool after prepare_tools has withdrawn it; the post-grace state gets the plain no-tools text back. - Split search-budget rejections from any failed tool call: the code tool raises ToolFailed for every error in model-written Python, so budget_spent was true for a ZeroDivisionError. - docs/capabilities/rag.md described the old single-turn removal. - Drop the rationale clause from the CHANGELOG entry. --- CHANGELOG.md | 4 +- docs/capabilities/rag.md | 2 +- evaluations/evaluations/benchmark.py | 3 +- evaluations/evaluations/capability_runner.py | 39 +++++++++++-------- evaluations/tests/test_capability_runner.py | 32 ++++++++++++++- .../haiku/rag/capabilities/_base.py | 13 ++++++- tests/capabilities/test_capabilities.py | 9 +++++ 7 files changed, 80 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 838472a8..3e25dbba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ ### Added -- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_calls`, `n_executions`, `n_requests` and `budget_spent` as eval attributes alongside `cited_uris`. +- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_searches`, `n_failed_tools`, `n_executions`, `n_requests` and `budget_spent` as eval attributes alongside `cited_uris`. ### Changed @@ -12,7 +12,7 @@ ### Fixed - A failed `analysis_execute_code` call that iterated a file object reports the `.readlines()` workaround alongside the `TypeError`. -- A capability that reaches its request limit keeps its cite tool for two further model requests while its other tools are removed, so an exhausted run can still register citations. +- A capability that reaches its request limit keeps its cite tool for two further model requests while its other tools are removed. - A cited chunk id that does not match a retrieved id exactly resolves to the closest one above a 0.75 similarity cutoff, recovering ids mistyped from search results. - Dotfiles are parsed as their actual format instead of a single unstructured text block. Docling ignores the extension of a name starting with a dot, so converters strip leading dots from the name they hand it. diff --git a/docs/capabilities/rag.md b/docs/capabilities/rag.md index af3b2959..6dce2443 100644 --- a/docs/capabilities/rag.md +++ b/docs/capabilities/rag.md @@ -26,7 +26,7 @@ print(result.output) `create_capability` accepts `db_path`, `config`, `defer_loading`, `request_limit`, and `vision`. Set `defer_loading=False` for a dedicated RAG agent where routing is unnecessary. The default request limit is 20 model requests per question; set `request_limit=None` to disable it. `vision` controls whether picture results are attached to search returns as images and should reflect the model the hosting agent runs; it defaults to the configured QA model's `vision` flag. -When the limit is reached, only the RAG capability's tools are removed. The model gets one more turn to answer from evidence already gathered, while unrelated agent and capability tools remain available. A new agent run starts a fresh limit, so multi-turn chat does not consume one shared budget. +When the limit is reached, `rag_search` is removed while `rag_cite` remains for two further requests, so the model can register citations before answering from evidence already gathered. Unrelated agent and capability tools remain available. A new agent run starts a fresh limit, so multi-turn chat does not consume one shared budget. ## State diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index f2339b82..51cef761 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -438,7 +438,8 @@ async def run_qa_benchmark( set_eval_attribute("searched_uris", result.searched_uris) set_eval_attribute("n_searches", result.n_searches) set_eval_attribute("n_search_calls", result.n_search_calls) - set_eval_attribute("n_rejected_calls", result.n_rejected_calls) + set_eval_attribute("n_rejected_searches", result.n_rejected_searches) + set_eval_attribute("n_failed_tools", result.n_failed_tools) set_eval_attribute("n_executions", result.n_executions) set_eval_attribute("n_requests", result.n_requests) set_eval_attribute("budget_spent", result.budget_spent) diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index d5a4203c..03da8773 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -36,24 +36,30 @@ class CapabilityRunResult: n_searches: int = 0 n_executions: int = 0 n_search_calls: int = 0 - n_rejected_calls: int = 0 + n_rejected_searches: int = 0 + n_failed_tools: int = 0 n_requests: int = 0 budget_spent: bool = False def _count_tool_traffic( messages: list[ModelMessage], namespace: str -) -> tuple[int, int, int]: - """Count search calls, rejected calls and model requests in a run. +) -> tuple[int, int, int, int]: + """Count search calls, failed calls and model requests in a run. ``state.searches`` is keyed by query, so it collapses repeated queries and never records a call the capability refused. Counting the message history - instead gives the real number of attempts, which is what shows whether a - case ran out of budget. + instead gives the real number of attempts. + + Failures are split by tool. Only the search tool fails for want of budget, + whereas the code tool raises ``ToolFailed`` for any error in model-written + Python, so counting every failure together would report a ``ZeroDivisionError`` + as budget exhaustion. """ search_tool = f"{namespace}_search" search_calls = 0 - rejected = 0 + rejected_searches = 0 + failed_tools = 0 requests = 0 for message in messages: if isinstance(message, ModelResponse): @@ -64,12 +70,12 @@ def _count_tool_traffic( if isinstance(part, ToolCallPart) and part.tool_name == search_tool ) continue - rejected += sum( - 1 - for part in message.parts - if isinstance(part, ToolReturnPart) and part.outcome == "failed" - ) - return search_calls, rejected, requests + for part in message.parts: + if isinstance(part, ToolReturnPart) and part.outcome == "failed": + failed_tools += 1 + if part.tool_name == search_tool: + rejected_searches += 1 + return search_calls, rejected_searches, failed_tools, requests @dataclass @@ -141,8 +147,8 @@ async def run_capability_question( executions = getattr(state, "executions", None) n_executions = len(executions) if executions is not None else 0 - n_search_calls, n_rejected_calls, n_requests = _count_tool_traffic( - agent_result.all_messages(), capability.state_namespace + n_search_calls, n_rejected_searches, n_failed_tools, n_requests = ( + _count_tool_traffic(agent_result.all_messages(), capability.state_namespace) ) return CapabilityRunResult( @@ -153,7 +159,8 @@ async def run_capability_question( n_searches=len(typed.searches), n_executions=n_executions, n_search_calls=n_search_calls, - n_rejected_calls=n_rejected_calls, + n_rejected_searches=n_rejected_searches, + n_failed_tools=n_failed_tools, n_requests=n_requests, - budget_spent=n_rejected_calls > 0, + budget_spent=n_rejected_searches > 0, ) diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index cd1b0dea..8e3cf317 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/evaluations/tests/test_capability_runner.py @@ -18,6 +18,34 @@ from haiku.rag.capabilities.rag import create_capability as create_rag from haiku.rag.config.models import AppConfig +def test_count_tool_traffic_separates_search_rejections_from_code_errors(): + """A crash in model-written Python must not read as budget exhaustion.""" + messages = [ + ModelRequest(parts=[UserPromptPart(content="q")]), + ModelResponse(parts=[ToolCallPart("analysis_execute_code", {"code": "1/0"})]), + ModelRequest( + parts=[ + ToolReturnPart( + tool_name="analysis_execute_code", + content="ZeroDivisionError", + tool_call_id="1", + outcome="failed", + ) + ] + ), + ModelResponse(parts=[TextPart("done")]), + ] + + search_calls, rejected_searches, failed_tools, requests = _count_tool_traffic( + messages, "analysis" + ) + + assert search_calls == 0 + assert rejected_searches == 0 + assert failed_tools == 1 + assert requests == 2 + + def test_count_tool_traffic_counts_attempts_not_distinct_queries(): """Rejected and repeated calls both count; `state.searches` hides them.""" messages = [ @@ -44,7 +72,9 @@ def test_count_tool_traffic_counts_attempts_not_distinct_queries(): ModelResponse(parts=[TextPart("done")]), ] - search_calls, rejected, requests = _count_tool_traffic(messages, "analysis") + search_calls, rejected, _failed, requests = _count_tool_traffic( + messages, "analysis" + ) assert search_calls == 2 assert rejected == 1 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 66e09efd..9bd51bf9 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -171,7 +171,18 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): return request_context def _budget_notice(self) -> str | None: - """Tell the model which of this capability's budgets just ran out.""" + """Tell the model which of this capability's budgets just ran out. + + Never names a tool ``prepare_tools`` has already withdrawn: pointing the + model at a tool that is gone costs it the agent's unknown-tool retry + budget and can abort the run. + """ + if self._citation_grace_expired: + return ( + f"The {self.state_namespace} capability's tools are no longer " + "available. Give the best answer possible using the evidence " + "already gathered." + ) if self._request_limit_reached: return ( f"The {self.state_namespace} capability has reached its request " diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index ac99b1d9..e062bd9c 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -583,10 +583,19 @@ async def test_cite_tool_is_withdrawn_after_the_grace_window(temp_db_path): capability.request_count = 2 kept = await capability.prepare_tools(ctx, cast(Any, tool_defs)) assert {tool.name for tool in kept} == {"rag_cite"} + notice = capability._budget_notice() + assert notice is not None and "rag_cite" in notice capability.request_count = 4 kept = await capability.prepare_tools(ctx, cast(Any, tool_defs)) assert kept == [] + # The notice must never point at a tool prepare_tools has withdrawn: + # calling a missing tool burns the agent's unknown-tool retries and can + # abort the run. + notice = capability._budget_notice() + assert notice is not None + assert "rag_cite" not in notice + assert "no longer available" in notice @pytest.mark.asyncio