trace the output tool call but not count it

Author: alibeyram

pydantic_ai_slim/pydantic_ai/_tool_manager.py

@@ -107,16 +107,7 @@ async def handle_call(
         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
 
-        if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
-            # Output tool calls are not traced and not counted
-            return await self._call_tool(
-                call,
-                allow_partial=allow_partial,
-                wrap_validation_errors=wrap_validation_errors,
-                approved=approved,
-            )
-        else:
-            return await self._call_function_tool(
+        return await self._call_function_tool(
                 call,
                 allow_partial=allow_partial,
                 wrap_validation_errors=wrap_validation_errors,
@@ -217,12 +208,19 @@ async def _call_function_tool(
         """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
         instrumentation_names = InstrumentationNames.for_version(instrumentation_version)
 
+        if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
+            tool_name: str = 'Output Tool'
+            output_tool_flag = True
+        else:
+            tool_name: str = call.tool_name
+            output_tool_flag = False
+
         span_attributes = {
-            'gen_ai.tool.name': call.tool_name,
+            'gen_ai.tool.name': tool_name,
             # NOTE: this means `gen_ai.tool.call.id` will be included even if it was generated by pydantic-ai
             'gen_ai.tool.call.id': call.tool_call_id,
             **({instrumentation_names.tool_arguments_attr: call.args_as_json_str()} if include_content else {}),
-            'logfire.msg': f'running tool: {call.tool_name}',
+            'logfire.msg': f'running tool: {tool_name}',
             # add the JSON schema so these attributes are formatted nicely in Logfire
             'logfire.json_schema': json.dumps(
                 {
@@ -243,7 +241,7 @@ async def _call_function_tool(
             ),
         }
         with tracer.start_as_current_span(
-            instrumentation_names.get_tool_span_name(call.tool_name),
+            instrumentation_names.get_tool_span_name(tool_name),
             attributes=span_attributes,
         ) as span:
             try:
@@ -253,7 +251,9 @@ async def _call_function_tool(
                     wrap_validation_errors=wrap_validation_errors,
                     approved=approved,
                 )
-                usage.tool_calls += 1
+                if not output_tool_flag:
+                    # Output tool calls are not counted
+                    usage.tool_calls += 1
 
             except ToolRetryError as e:
                 part = e.tool_retry