Record tool response in tool run span (#2109)

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/tools.py

@@ -32,6 +32,7 @@
     'ToolDefinition',
 )
 
+from .messages import ToolReturnPart
 
 ToolParams = ParamSpec('ToolParams', default=...)
 """Retrieval function param spec."""
@@ -346,15 +347,31 @@ async def run(
                 {
                     'type': 'object',
                     'properties': {
-                        **({'tool_arguments': {'type': 'object'}} if include_content else {}),
+                        **(
+                            {
+                                'tool_arguments': {'type': 'object'},
+                                'tool_response': {'type': 'object'},
+                            }
+                            if include_content
+                            else {}
+                        ),
                         'gen_ai.tool.name': {},
                         'gen_ai.tool.call.id': {},
                     },
                 }
             ),
         }
-        with tracer.start_as_current_span('running tool', attributes=span_attributes):
-            return await self._run(message, run_context)
+        with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
+            response = await self._run(message, run_context)
+            if include_content and span.is_recording():
+                span.set_attribute(
+                    'tool_response',
+                    response.model_response_str()
+                    if isinstance(response, ToolReturnPart)
+                    else response.model_response(),
+                )
+
+            return response
 
     async def _run(
         self, message: _messages.ToolCallPart, run_context: RunContext[AgentDepsT]

tests/test_logfire.py

@@ -554,6 +554,46 @@ async def add_numbers(x: int, y: int) -> int:
     ]
 
     if include_content:
-        assert tool_attributes['tool_arguments'] == snapshot('{"x":42,"y":42}')
+        assert tool_attributes == snapshot(
+            {
+                'gen_ai.tool.name': 'add_numbers',
+                'gen_ai.tool.call.id': IsStr(),
+                'tool_arguments': '{"x":42,"y":42}',
+                'tool_response': '84',
+                'logfire.msg': 'running tool: add_numbers',
+                'logfire.json_schema': IsJson(
+                    snapshot(
+                        {
+                            'type': 'object',
+                            'properties': {
+                                'tool_arguments': {'type': 'object'},
+                                'tool_response': {'type': 'object'},
+                                'gen_ai.tool.name': {},
+                                'gen_ai.tool.call.id': {},
+                            },
+                        }
+                    )
+                ),
+                'logfire.span_type': 'span',
+            }
+        )
     else:
-        assert 'tool_arguments' not in tool_attributes
+        assert tool_attributes == snapshot(
+            {
+                'gen_ai.tool.name': 'add_numbers',
+                'gen_ai.tool.call.id': IsStr(),
+                'logfire.msg': 'running tool: add_numbers',
+                'logfire.json_schema': IsJson(
+                    snapshot(
+                        {
+                            'type': 'object',
+                            'properties': {
+                                'gen_ai.tool.name': {},
+                                'gen_ai.tool.call.id': {},
+                            },
+                        }
+                    )
+                ),
+                'logfire.span_type': 'span',
+            }
+        )