refactor: use presence of usage_limits parameter to determine tool counting

Author: tradeqvest

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -701,7 +701,7 @@ async def process_function_tools(  # noqa: C901
             output_parts.append(part)
         else:
             try:
-                result_data = await tool_manager.handle_call(call, usage_limits=ctx.deps.usage_limits)
+                result_data = await tool_manager.handle_call(call)
             except exceptions.UnexpectedModelBehavior as e:
                 ctx.state.increment_retries(ctx.deps.max_result_retries, e)
                 raise e  # pragma: lax no cover
@@ -764,9 +764,9 @@ async def process_function_tools(  # noqa: C901
             calls_to_run,
             deferred_tool_results,
             ctx.deps.tracer,
+            ctx.deps.usage_limits,
             output_parts,
             deferred_calls,
-            ctx.deps.usage_limits,
         ):
             yield event
 
@@ -811,9 +811,9 @@ async def _call_tools(
     tool_calls: list[_messages.ToolCallPart],
     deferred_tool_results: dict[str, DeferredToolResult],
     tracer: Tracer,
+    usage_limits: _usage.UsageLimits | None,
     output_parts: list[_messages.ModelRequestPart],
     output_deferred_calls: dict[Literal['external', 'unapproved'], list[_messages.ToolCallPart]],
-    usage_limits: _usage.UsageLimits | None,
 ) -> AsyncIterator[_messages.HandleResponseEvent]:
     tool_parts_by_index: dict[int, _messages.ModelRequestPart] = {}
     user_parts_by_index: dict[int, _messages.UserPromptPart] = {}

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -86,7 +86,7 @@ async def handle_call(
 
         if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
             # Output tool calls are not traced
-            return await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits, is_output_tool=True)
+            return await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
         else:
             return await self._call_tool_traced(
                 call,
@@ -103,7 +103,6 @@ async def _call_tool(
         allow_partial: bool,
         wrap_validation_errors: bool,
         usage_limits: UsageLimits | None,
-        is_output_tool: bool = False,
     ) -> Any:
         if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
@@ -135,12 +134,12 @@ async def _call_tool(
             else:
                 args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
 
-            if not is_output_tool and usage_limits is not None:
+            if usage_limits is not None:
                 usage_limits.check_before_tool_call(self.ctx.usage)
 
             result = await self.toolset.call_tool(name, args_dict, ctx, tool)
 
-            if not is_output_tool:
+            if usage_limits is not None:
                 self.ctx.usage.tool_calls += 1
 
             return result