From f1cd425d1e836273c21887ce5335cfdf6f881b30 Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 14:17:02 -0800 Subject: [PATCH 1/7] trace the output tool call but not count it --- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index fb7039e2cc..1f27f3a4bd 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/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 .""" 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 From 419d5ace6b8caf938de5dd0baea680d9ca197c0f Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 14:37:19 -0800 Subject: [PATCH 2/7] some fixes --- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index 1f27f3a4bd..5c997749c4 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -106,6 +106,11 @@ 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_flag = True + else: + output_tool_flag = False return await self._call_function_tool( call, @@ -116,6 +121,7 @@ async def handle_call( include_content=self.ctx.trace_include_content, instrumentation_version=self.ctx.instrumentation_version, usage=self.ctx.usage, + output_tool_flag=output_tool_flag, ) async def _call_tool( @@ -204,16 +210,17 @@ async def _call_function_tool( include_content: bool, instrumentation_version: int, usage: RunUsage, + output_tool_flag: bool = False, ) -> Any: """See .""" 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 + if output_tool_flag: + tool_name = 'output tool' else: - tool_name: str = call.tool_name - output_tool_flag = False + tool_name = call.tool_name + + span_attributes = { 'gen_ai.tool.name': tool_name, From 0b287bd4a7830988dbde71fa53888d89660a87fd Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 14:54:38 -0800 Subject: [PATCH 3/7] linter --- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 24 +++++++++---------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index 5c997749c4..dde27df1e6 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -106,23 +106,23 @@ 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_flag = True else: output_tool_flag = False return await self._call_function_tool( - call, - allow_partial=allow_partial, - wrap_validation_errors=wrap_validation_errors, - approved=approved, - tracer=self.ctx.tracer, - include_content=self.ctx.trace_include_content, - instrumentation_version=self.ctx.instrumentation_version, - usage=self.ctx.usage, - output_tool_flag=output_tool_flag, - ) + call, + allow_partial=allow_partial, + wrap_validation_errors=wrap_validation_errors, + approved=approved, + tracer=self.ctx.tracer, + include_content=self.ctx.trace_include_content, + instrumentation_version=self.ctx.instrumentation_version, + usage=self.ctx.usage, + output_tool_flag=output_tool_flag, + ) async def _call_tool( self, @@ -220,8 +220,6 @@ async def _call_function_tool( else: tool_name = call.tool_name - - span_attributes = { '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 From 11f6fe7840186455911305a53b86d53603488af6 Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 16:01:16 -0800 Subject: [PATCH 4/7] update the test logfire --- tests/test_logfire.py | 64 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 52 insertions(+), 12 deletions(-) diff --git a/tests/test_logfire.py b/tests/test_logfire.py index dadb930dd0..333d91e81e 100644 --- a/tests/test_logfire.py +++ b/tests/test_logfire.py @@ -692,7 +692,10 @@ class MyOutput: 'id': 0, 'name': 'agent run', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'running tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -703,7 +706,10 @@ class MyOutput: 'id': 0, 'name': 'invoke_agent my_agent', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'execute_tool output tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -900,7 +906,10 @@ class MyOutput: 'id': 0, 'name': 'agent run', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'running tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -911,7 +920,10 @@ class MyOutput: 'id': 0, 'name': 'invoke_agent my_agent', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'execute_tool output tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -1381,8 +1393,15 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, { 'id': 2, - 'name': 'running output function', - 'message': 'running output function: final_result', + 'name': 'running tool', + 'message': 'running tool: output tool', + 'children': [ + { + 'id': 3, + 'name': 'running output function', + 'message': 'running output function: final_result', + } + ], }, ], } @@ -1428,8 +1447,15 @@ def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse: {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'}, { 'id': 2, - 'name': 'execute_tool final_result', - 'message': 'running output function: final_result', + 'name': 'execute_tool output tool', + 'message': 'running tool: output tool', + 'children': [ + { + 'id': 3, + 'name': 'execute_tool final_result', + 'message': 'running output function: final_result', + } + ], }, ], } @@ -2336,7 +2362,10 @@ def instructions(): 'id': 0, 'name': 'agent run', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'running tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -2347,7 +2376,10 @@ def instructions(): 'id': 0, 'name': 'invoke_agent my_agent', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'execute_tool output tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -2589,6 +2621,7 @@ def my_tool() -> str: 'children': [{'id': 3, 'name': 'running tool', 'message': 'running tool: my_tool'}], }, {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + {'id': 5, 'name': 'running tool', 'message': 'running tool: output tool'}, ], } ] @@ -2611,6 +2644,7 @@ def my_tool() -> str: ], }, {'id': 4, 'name': 'chat test', 'message': 'chat test'}, + {'id': 5, 'name': 'execute_tool output tool', 'message': 'running tool: output tool'}, ], } ] @@ -2877,7 +2911,10 @@ def instructions(ctx: RunContext[None]): 'id': 0, 'name': 'agent run', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'running tool', 'message': 'running tool: output tool'}, + ], } ] ) @@ -2888,7 +2925,10 @@ def instructions(ctx: RunContext[None]): 'id': 0, 'name': 'invoke_agent my_agent', 'message': 'my_agent run', - 'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}], + 'children': [ + {'id': 1, 'name': 'chat test', 'message': 'chat test'}, + {'id': 2, 'name': 'execute_tool output tool', 'message': 'running tool: output tool'}, + ], } ] ) From b8ad3df5914c210284cc0df7a5f5addd0df2d14f Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 16:06:35 -0800 Subject: [PATCH 5/7] update test temporal --- tests/test_temporal.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_temporal.py b/tests/test_temporal.py index b3fe75d911..5d0c49ea44 100644 --- a/tests/test_temporal.py +++ b/tests/test_temporal.py @@ -718,6 +718,7 @@ async def test_complex_agent_run_in_workflow( ) ], ), + BasicSpan(content='running tool: output tool'), ], ), BasicSpan(content='CompleteWorkflow:ComplexAgentWorkflow'), From 7bbb7e3b9ebaa55bc90c42b529f01a02cd1c4419 Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 16:29:28 -0800 Subject: [PATCH 6/7] test prefect --- tests/test_prefect.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_prefect.py b/tests/test_prefect.py index 3f2059ab6f..3ab9422364 100644 --- a/tests/test_prefect.py +++ b/tests/test_prefect.py @@ -553,6 +553,7 @@ async def run_complex_agent() -> Response: ) ], ), + BasicSpan(content='running tool: output tool'), ], ) ], From e26e10dfab2fb404ad55cb9016b8d72d9bb356b1 Mon Sep 17 00:00:00 2001 From: alibeyram Date: Tue, 11 Nov 2025 16:39:32 -0800 Subject: [PATCH 7/7] test dbos --- tests/test_dbos.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_dbos.py b/tests/test_dbos.py index 256aba83fb..d1ea9f8706 100644 --- a/tests/test_dbos.py +++ b/tests/test_dbos.py @@ -594,6 +594,7 @@ async def test_complex_agent_run_in_workflow(allow_model_requests: None, dbos: D ) ], ), + BasicSpan(content='running tool: output tool'), ], ) ],