Skip to content

Commit 6018ae1

Browse files
committed
refactor: remove usage_limits parameter from _call_tool
- Simplified the tool call logic by removing the unused usage_limits parameter from the _call_tool method in ToolManager.
1 parent 80abb7b commit 6018ae1

File tree

2 files changed

+35
-2
lines changed

2 files changed

+35
-2
lines changed

pydantic_ai_slim/pydantic_ai/_tool_manager.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,6 @@ async def _call_tool(
125125
call: ToolCallPart,
126126
allow_partial: bool,
127127
wrap_validation_errors: bool,
128-
usage_limits: UsageLimits | None = None,
129128
count_tool_usage: bool = True,
130129
) -> Any:
131130
if self.tools is None or self.ctx is None:
@@ -239,7 +238,7 @@ async def _call_tool_traced(
239238
attributes=span_attributes,
240239
) as span:
241240
try:
242-
tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
241+
tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
243242
except ToolRetryError as e:
244243
part = e.tool_retry
245244
if include_content and span.is_recording():

tests/test_usage_limits.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,40 @@ async def another_regular_tool(x: str) -> str:
290290
assert result_output.usage() == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=15, tool_calls=1))
291291

292292

293+
async def test_output_tool_allowed_at_limit() -> None:
294+
"""Test that output tools can be called even when at the tool_calls_limit."""
295+
296+
class MyOutput(BaseModel):
297+
result: str
298+
299+
def call_output_after_regular(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
300+
if len(messages) == 1:
301+
return ModelResponse(
302+
parts=[
303+
ToolCallPart('regular_tool', {'x': 'test'}, 'call_1'),
304+
],
305+
usage=RequestUsage(input_tokens=10, output_tokens=5),
306+
)
307+
else:
308+
return ModelResponse(
309+
parts=[
310+
ToolCallPart('final_result', {'result': 'success'}, 'call_2'),
311+
],
312+
usage=RequestUsage(input_tokens=10, output_tokens=5),
313+
)
314+
315+
test_agent = Agent(FunctionModel(call_output_after_regular), output_type=ToolOutput(MyOutput))
316+
317+
@test_agent.tool_plain
318+
async def regular_tool(x: str) -> str:
319+
return f'{x}-processed'
320+
321+
result = await test_agent.run('test', usage_limits=UsageLimits(tool_calls_limit=1))
322+
323+
assert result.output.result == 'success'
324+
assert result.usage() == snapshot(RunUsage(requests=2, input_tokens=20, output_tokens=10, tool_calls=1))
325+
326+
293327
async def test_failed_tool_calls_not_counted() -> None:
294328
"""Test that failed tool calls (raising ModelRetry) are not counted."""
295329
test_agent = Agent(TestModel())

0 commit comments

Comments
 (0)