Scope the limit notice and spend the cite window on own turns only
The request-limit notice said only the cite tool remained available, but chat registers rag and analysis in one agent, so exhausting analysis claimed rag_search was gone too. Scoped to the capability's own tools. The cite window was counted over every model request once loaded, so turns spent on another capability expired it before the model was ever placed where citing was the obvious move. Count only requests whose preceding response called one of this capability's tools; engagement is also the only thing that can loop, which is all the bound guards against. Also: _count_tool_traffic returns a named tuple rather than four bare ints, and counts failures only for this capability's tools, so host-tool retries and output-validation retries no longer read as its failures.
This commit is contained in:
Yiorgis Gozadinos
parent
27ed0b3bb3
commit
c137305468
7 changed files with 148 additions and 79 deletions
|
|
@ -12,12 +12,13 @@
|
|||
### 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.
|
||||
- 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.
|
||||
- A capability that reaches its request limit keeps its cite tool for two further requests that call one of its tools, 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.
|
||||
- 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.
|
||||
|
||||
### Documentation
|
||||
|
||||
- `qa.max_searches` is documented as defaulting to 5 in `docs/configuration/qa.md` and `docs/configuration/index.md`, was 3.
|
||||
- New "Vacuum Memory Requirements" subsection in `docs/configuration/storage.md` documenting the ~5x peak memory of vacuum compaction relative to the `documents` table, the mitigations, and upstream issue lancedb/lancedb#2325.
|
||||
|
||||
## [0.71.0] - 2026-07-29
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@
|
|||
|
||||
It is deferred by default, keeping its substantial instructions and tool schemas out of context until the model chooses to load it.
|
||||
|
||||
The default request limit is 30 model requests per question. Override it with `create_capability(request_limit=...)`, or set `request_limit=None` to disable it. As with the RAG capability, `create_capability(vision=...)` overrides the image-attachment gate, defaulting to the configured analysis model's `vision` flag. At the limit, `analysis_search` and `analysis_execute_code` are removed while `analysis_cite` remains for two further requests, so the model can register citations before answering from gathered evidence. Other agent and capability tools remain available, and the budget resets for every agent run.
|
||||
The default request limit is 30 model requests per question. Override it with `create_capability(request_limit=...)`, or set `request_limit=None` to disable it. As with the RAG capability, `create_capability(vision=...)` overrides the image-attachment gate, defaulting to the configured analysis model's `vision` flag. At the limit, `analysis_search` and `analysis_execute_code` are removed while `analysis_cite` remains for two further requests that call an analysis tool, so the model can register citations before answering from gathered evidence. Requests spent on other capabilities do not count against that window. Other agent and capability tools remain available, and the budget resets for every agent run.
|
||||
|
||||
When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool keeps failing rather than disappearing, and the instructions name it on every following request. Searching from inside `analysis_execute_code` does not count against `qa.max_searches`.
|
||||
|
||||
|
|
|
|||
|
|
@ -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, `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.
|
||||
When the limit is reached, `rag_search` is removed while `rag_cite` remains for two further requests that call a RAG tool, so the model can register citations before answering from evidence already gathered. Requests spent on other capabilities do not count against that window. 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
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol, cast
|
||||
from typing import Any, NamedTuple, Protocol, cast
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import (
|
||||
|
|
@ -42,32 +42,32 @@ class CapabilityRunResult:
|
|||
n_requests: int = 0
|
||||
|
||||
|
||||
class ToolTraffic(NamedTuple):
|
||||
n_search_calls: int
|
||||
n_rejected_searches: int
|
||||
n_failed_tools: int
|
||||
n_requests: int
|
||||
|
||||
|
||||
def _count_tool_traffic(
|
||||
messages: list[ModelMessage], namespace: str
|
||||
) -> tuple[int, int, int, int]:
|
||||
messages: list[ModelMessage], namespace: str, tool_names: frozenset[str]
|
||||
) -> ToolTraffic:
|
||||
"""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. The capability object cannot be
|
||||
read instead: ``for_run`` hands the run a ``replace()`` copy, so the outer
|
||||
instance's counters stay at zero. Counting the message history is the only
|
||||
way to see the real number of attempts.
|
||||
The history is the only source: ``state.searches`` is keyed by query so it
|
||||
hides repeats and refusals, and ``for_run`` hands the run a ``replace()``
|
||||
copy, leaving the outer capability's counters at zero.
|
||||
|
||||
Search failures are counted apart because the search tool fails only when
|
||||
its budget is spent. A failed code call is ambiguous — an exhausted
|
||||
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.
|
||||
Only search failures mean an exhausted budget. A failed code call may be
|
||||
either the execution budget or any error in model-written Python, so
|
||||
``n_failed_tools`` covers both without claiming to tell them apart. It counts
|
||||
``RetryPromptPart`` too, since ``_cite`` rejects with ``ModelRetry`` and only
|
||||
``ToolFailed`` sets ``outcome="failed"``. Both are restricted to
|
||||
``tool_names``, excluding host tools and output-validation retries.
|
||||
|
||||
``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.
|
||||
``n_requests`` counts the run's requests, which matches the capability's own
|
||||
budget only while it stays loaded — a deferred capability skips hooks until
|
||||
it loads.
|
||||
"""
|
||||
search_tool = f"{namespace}_search"
|
||||
search_calls = 0
|
||||
|
|
@ -84,13 +84,22 @@ def _count_tool_traffic(
|
|||
)
|
||||
continue
|
||||
for part in message.parts:
|
||||
if not isinstance(part, RetryPromptPart | ToolReturnPart):
|
||||
continue
|
||||
if part.tool_name not in tool_names:
|
||||
continue
|
||||
if isinstance(part, RetryPromptPart):
|
||||
failed_tools += 1
|
||||
elif isinstance(part, ToolReturnPart) and part.outcome == "failed":
|
||||
elif part.outcome == "failed":
|
||||
failed_tools += 1
|
||||
if part.tool_name == search_tool:
|
||||
rejected_searches += 1
|
||||
return search_calls, rejected_searches, failed_tools, requests
|
||||
return ToolTraffic(
|
||||
n_search_calls=search_calls,
|
||||
n_rejected_searches=rejected_searches,
|
||||
n_failed_tools=failed_tools,
|
||||
n_requests=requests,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -162,8 +171,10 @@ 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_searches, n_failed_tools, n_requests = (
|
||||
_count_tool_traffic(agent_result.all_messages(), capability.state_namespace)
|
||||
traffic = _count_tool_traffic(
|
||||
agent_result.all_messages(),
|
||||
capability.state_namespace,
|
||||
capability.tool_names,
|
||||
)
|
||||
|
||||
return CapabilityRunResult(
|
||||
|
|
@ -171,10 +182,14 @@ async def run_capability_question(
|
|||
cited_uris=cited_uris,
|
||||
cited_chunk_ids=cited_chunk_ids,
|
||||
searched_uris=searched_uris,
|
||||
# Distinct search keys, not searches. Analysis files every in-code
|
||||
# `search()` under one "_sandbox" key, so twenty sandbox searches read
|
||||
# as one here; `n_search_calls` is the true count of search *tool*
|
||||
# calls, and in-code searches are not counted anywhere.
|
||||
n_searches=len(typed.searches),
|
||||
n_executions=n_executions,
|
||||
n_search_calls=n_search_calls,
|
||||
n_rejected_searches=n_rejected_searches,
|
||||
n_failed_tools=n_failed_tools,
|
||||
n_requests=n_requests,
|
||||
n_search_calls=traffic.n_search_calls,
|
||||
n_rejected_searches=traffic.n_rejected_searches,
|
||||
n_failed_tools=traffic.n_failed_tools,
|
||||
n_requests=traffic.n_requests,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,10 @@ from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
|||
from haiku.rag.capabilities.rag import create_capability as create_rag
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
ANALYSIS_TOOLS = frozenset(
|
||||
{"analysis_search", "analysis_execute_code", "analysis_cite"}
|
||||
)
|
||||
|
||||
|
||||
def test_count_tool_traffic_sees_a_rejected_cite_call():
|
||||
"""`_cite` rejects with ModelRetry, which is not a failed ToolReturnPart."""
|
||||
|
|
@ -36,12 +40,10 @@ def test_count_tool_traffic_sees_a_rejected_cite_call():
|
|||
ModelResponse(parts=[TextPart("done")]),
|
||||
]
|
||||
|
||||
_search_calls, rejected_searches, failed_tools, _requests = _count_tool_traffic(
|
||||
messages, "analysis"
|
||||
)
|
||||
traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
|
||||
|
||||
assert failed_tools == 1
|
||||
assert rejected_searches == 0
|
||||
assert traffic.n_failed_tools == 1
|
||||
assert traffic.n_rejected_searches == 0
|
||||
|
||||
|
||||
def test_count_tool_traffic_separates_search_rejections_from_code_errors():
|
||||
|
|
@ -62,14 +64,12 @@ def test_count_tool_traffic_separates_search_rejections_from_code_errors():
|
|||
ModelResponse(parts=[TextPart("done")]),
|
||||
]
|
||||
|
||||
search_calls, rejected_searches, failed_tools, requests = _count_tool_traffic(
|
||||
messages, "analysis"
|
||||
)
|
||||
traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
|
||||
|
||||
assert search_calls == 0
|
||||
assert rejected_searches == 0
|
||||
assert failed_tools == 1
|
||||
assert requests == 2
|
||||
assert traffic.n_search_calls == 0
|
||||
assert traffic.n_rejected_searches == 0
|
||||
assert traffic.n_failed_tools == 1
|
||||
assert traffic.n_requests == 2
|
||||
|
||||
|
||||
def test_count_tool_traffic_counts_attempts_not_distinct_queries():
|
||||
|
|
@ -98,13 +98,11 @@ def test_count_tool_traffic_counts_attempts_not_distinct_queries():
|
|||
ModelResponse(parts=[TextPart("done")]),
|
||||
]
|
||||
|
||||
search_calls, rejected, _failed, requests = _count_tool_traffic(
|
||||
messages, "analysis"
|
||||
)
|
||||
traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
|
||||
|
||||
assert search_calls == 2
|
||||
assert rejected == 1
|
||||
assert requests == 2
|
||||
assert traffic.n_search_calls == 2
|
||||
assert traffic.n_rejected_searches == 1
|
||||
assert traffic.n_requests == 2
|
||||
|
||||
|
||||
async def test_runs_rag_capability_without_legacy_capability_layer(tmp_path):
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@ from pydantic_ai.messages import (
|
|||
InstructionPart,
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
ToolCallPart,
|
||||
ToolReturn,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
|
|
@ -29,7 +31,13 @@ from haiku.rag.store.models.citation import Citation, resolve_citations
|
|||
from haiku.rag.tools.search import build_binary_parts_from_results
|
||||
|
||||
CITATION_GRACE_REQUESTS = 2
|
||||
"""Requests the cite tool outlives this capability's other tools by."""
|
||||
"""Requests calling this capability's tools that its cite tool outlives the rest by.
|
||||
|
||||
A loop guard, not a budget: cite consumes no retry budget and raises nothing, so
|
||||
left available forever a stuck model calls it until the agent's own request limit
|
||||
raises ``UsageLimitExceeded`` and the question returns no answer at all. Only
|
||||
engagement can loop, which is why other capabilities' turns do not spend it.
|
||||
"""
|
||||
|
||||
CHUNK_ID_MATCH_CUTOFF = 0.75
|
||||
"""Similarity a cited chunk id needs to be treated as a corrupted known id.
|
||||
|
|
@ -105,6 +113,17 @@ def _compact_old_tool_returns(
|
|||
return compacted
|
||||
|
||||
|
||||
def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool:
|
||||
"""Whether the model's most recent response called one of these tools."""
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, ModelResponse):
|
||||
return any(
|
||||
isinstance(part, ToolCallPart) and part.tool_name in tool_names
|
||||
for part in message.parts
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
@dataclass
|
||||
class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
||||
db_path: Path
|
||||
|
|
@ -122,6 +141,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
|
||||
search_count: int = field(default=0, repr=False)
|
||||
request_count: int = field(default=0, repr=False)
|
||||
grace_requests_used: int = field(default=0, repr=False)
|
||||
|
||||
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
|
||||
outer = getattr(ctx.deps, "state", None)
|
||||
|
|
@ -138,6 +158,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
resource_lock=asyncio.Lock(),
|
||||
search_count=0,
|
||||
request_count=0,
|
||||
grace_requests_used=0,
|
||||
)
|
||||
run_capability._sync_state()
|
||||
return run_capability
|
||||
|
|
@ -167,6 +188,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
InstructionPart(content=instruction, dynamic=True),
|
||||
],
|
||||
)
|
||||
if self._request_limit_reached and _called_own_tool(
|
||||
request_context.messages, self.tool_names
|
||||
):
|
||||
self.grace_requests_used += 1
|
||||
self.request_count += 1
|
||||
return request_context
|
||||
|
||||
|
|
@ -186,9 +211,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
if self._request_limit_reached:
|
||||
return (
|
||||
f"The {self.state_namespace} capability has reached its request "
|
||||
f"limit. Only {self._cite_tool_name} remains available: register "
|
||||
"the chunk_ids supporting your answer, then answer from the "
|
||||
"evidence already gathered."
|
||||
f"limit. Only {self._cite_tool_name} remains among its tools: "
|
||||
"register the chunk_ids supporting your answer, then answer from "
|
||||
"the evidence already gathered."
|
||||
)
|
||||
if spent := self._spent_tool_names():
|
||||
names = ", ".join(sorted(spent))
|
||||
|
|
@ -215,14 +240,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
) -> list[ToolDefinition]:
|
||||
"""Remove this capability's tools past its limit, cite tool last.
|
||||
|
||||
The cite tool outlives the others by ``CITATION_GRACE_REQUESTS`` so a run
|
||||
that exhausts its budget can still record the evidence it gathered.
|
||||
|
||||
Tools whose own budget is spent stay declared on purpose. Removing one
|
||||
makes a model that calls it anyway hit ``Unknown tool name``, which is
|
||||
charged against the agent's unknown-tool retry budget and kills the run
|
||||
after two attempts. A spent tool that keeps failing wastes requests; a
|
||||
withdrawn one loses the whole answer.
|
||||
makes a model that calls it anyway hit ``Unknown tool name``, charged
|
||||
against the agent's unknown-tool retry budget, which kills the run after
|
||||
two attempts. A spent tool that keeps failing only wastes requests.
|
||||
"""
|
||||
if self._citation_grace_expired:
|
||||
return [tool for tool in tool_defs if tool.capability_id != self.id]
|
||||
|
|
@ -235,11 +256,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
]
|
||||
|
||||
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.
|
||||
"""
|
||||
"""Tools that can bring new evidence into the run."""
|
||||
return {f"{self.state_namespace}_search"}
|
||||
|
||||
def _spent_tool_names(self) -> set[str]:
|
||||
|
|
@ -264,10 +281,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
|
||||
@property
|
||||
def _citation_grace_expired(self) -> bool:
|
||||
return (
|
||||
self.request_limit is not None
|
||||
and self.request_count >= self.request_limit + CITATION_GRACE_REQUESTS
|
||||
)
|
||||
# No `request_limit is None` guard: the counter only advances under
|
||||
# `_request_limit_reached`, which already requires a limit.
|
||||
return self.grace_requests_used >= CITATION_GRACE_REQUESTS
|
||||
|
||||
async def after_run(
|
||||
self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
|
||||
|
|
|
|||
|
|
@ -17,7 +17,11 @@ from pydantic_ai.models.function import FunctionModel
|
|||
from pydantic_ai.models.test import TestModel
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from haiku.rag.capabilities._base import _compact_old_tool_returns
|
||||
from haiku.rag.capabilities._base import (
|
||||
CITATION_GRACE_REQUESTS,
|
||||
_called_own_tool,
|
||||
_compact_old_tool_returns,
|
||||
)
|
||||
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
|
||||
from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
||||
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState
|
||||
|
|
@ -494,6 +498,40 @@ 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]
|
||||
|
||||
|
||||
def test_grace_window_ignores_other_capabilities_turns():
|
||||
"""Only this capability's own tool calls may spend its cite window.
|
||||
|
||||
A multi-capability agent spends turns elsewhere; those must not expire the
|
||||
window that exists to give this capability a chance to cite.
|
||||
"""
|
||||
rag_tools = frozenset({"rag_search", "rag_cite"})
|
||||
|
||||
# Nothing to attribute before the model has responded at all.
|
||||
assert not _called_own_tool(
|
||||
[ModelRequest(parts=[UserPromptPart(content="q")])], rag_tools
|
||||
)
|
||||
assert not _called_own_tool(
|
||||
[ModelResponse(parts=[ToolCallPart("analysis_search", {"query": "x"})])],
|
||||
rag_tools,
|
||||
)
|
||||
assert not _called_own_tool(
|
||||
[ModelResponse(parts=[TextPart("just talking")])], rag_tools
|
||||
)
|
||||
assert _called_own_tool(
|
||||
[ModelResponse(parts=[ToolCallPart("rag_cite", {"chunk_ids": ["a"]})])],
|
||||
rag_tools,
|
||||
)
|
||||
# Only the most recent response counts, not any earlier one.
|
||||
assert not _called_own_tool(
|
||||
[
|
||||
ModelResponse(parts=[ToolCallPart("rag_cite", {"chunk_ids": ["a"]})]),
|
||||
ModelRequest(parts=[UserPromptPart(content="next")]),
|
||||
ModelResponse(parts=[ToolCallPart("analysis_search", {"query": "x"})]),
|
||||
],
|
||||
rag_tools,
|
||||
)
|
||||
|
||||
|
||||
@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.
|
||||
|
|
@ -511,13 +549,13 @@ async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_pa
|
|||
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.
|
||||
# Once the code budget is gone too there is nowhere left to send it.
|
||||
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
|
||||
assert "analysis_execute_code" in notice
|
||||
assert capability._evidence_tool_names() <= capability._spent_tool_names()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -531,7 +569,8 @@ async def test_spent_search_notice_tells_rag_to_answer(temp_db_path):
|
|||
notice = capability._budget_notice()
|
||||
|
||||
assert notice is not None
|
||||
assert "Answer from the evidence already gathered" in notice
|
||||
assert "rag_search" in notice
|
||||
assert capability._evidence_tool_names() == {"rag_search"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -591,7 +630,7 @@ async def test_exhausted_run_can_still_register_citations(rag_db):
|
|||
deps = Deps()
|
||||
|
||||
async with agent.iter("question", deps=deps) as run:
|
||||
async for node in run:
|
||||
async for _node in run:
|
||||
if chunk_id is None:
|
||||
searches = deps.state.get("rag", {}).get("searches") or {}
|
||||
for results in searches.values():
|
||||
|
|
@ -626,7 +665,7 @@ async def test_cite_tool_is_withdrawn_after_the_grace_window(temp_db_path):
|
|||
notice = capability._budget_notice()
|
||||
assert notice is not None and "rag_cite" in notice
|
||||
|
||||
capability.request_count = 4
|
||||
capability.grace_requests_used = CITATION_GRACE_REQUESTS
|
||||
kept = await capability.prepare_tools(ctx, cast(Any, tool_defs))
|
||||
assert kept == []
|
||||
# The notice must never point at a tool prepare_tools has withdrawn:
|
||||
|
|
|
|||
Loading…
Reference in a new issue