pydantic-ai-slim 0.0.30__py3-none-any.whl → 0.0.31__py3-none-any.whl

pydantic_ai/__init__.py

@@ -1,6 +1,6 @@
 from importlib.metadata import version
 
-from .agent import Agent, EndStrategy, HandleResponseNode, ModelRequestNode, UserPromptNode, capture_run_messages
+from .agent import Agent, CallToolsNode, EndStrategy, ModelRequestNode, UserPromptNode, capture_run_messages
 from .exceptions import (
     AgentRunError,
     FallbackExceptionGroup,
@@ -18,7 +18,7 @@ __all__ = (
     # agent
     'Agent',
     'EndStrategy',
-    'HandleResponseNode',
+    'CallToolsNode',
     'ModelRequestNode',
     'UserPromptNode',
     'capture_run_messages',

pydantic_ai/_agent_graph.py

@@ -23,6 +23,7 @@ from . import (
     result,
     usage as _usage,
 )
+from .models.instrumented import InstrumentedModel
 from .result import ResultDataT
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -36,7 +37,7 @@ __all__ = (
     'GraphAgentDeps',
     'UserPromptNode',
     'ModelRequestNode',
-    'HandleResponseNode',
+    'CallToolsNode',
     'build_run_context',
     'capture_run_messages',
 )
@@ -243,12 +244,12 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
 
     request: _messages.ModelRequest
 
-    _result: HandleResponseNode[DepsT, NodeRunEndT] | None = field(default=None, repr=False)
+    _result: CallToolsNode[DepsT, NodeRunEndT] | None = field(default=None, repr=False)
     _did_stream: bool = field(default=False, repr=False)
 
     async def run(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> HandleResponseNode[DepsT, NodeRunEndT]:
+    ) -> CallToolsNode[DepsT, NodeRunEndT]:
         if self._result is not None:
             return self._result
 
@@ -286,39 +287,33 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         assert not self._did_stream, 'stream() should only be called once per node'
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
-        with _logfire.span('model request', run_step=ctx.state.run_step) as span:
-            async with ctx.deps.model.request_stream(
-                ctx.state.message_history, model_settings, model_request_parameters
-            ) as streamed_response:
-                self._did_stream = True
-                ctx.state.usage.incr(_usage.Usage(), requests=1)
-                yield streamed_response
-                # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
-                # otherwise usage won't be properly counted:
-                async for _ in streamed_response:
-                    pass
-            model_response = streamed_response.get()
-            request_usage = streamed_response.usage()
-            span.set_attribute('response', model_response)
-            span.set_attribute('usage', request_usage)
+        async with ctx.deps.model.request_stream(
+            ctx.state.message_history, model_settings, model_request_parameters
+        ) as streamed_response:
+            self._did_stream = True
+            ctx.state.usage.incr(_usage.Usage(), requests=1)
+            yield streamed_response
+            # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
+            # otherwise usage won't be properly counted:
+            async for _ in streamed_response:
+                pass
+        model_response = streamed_response.get()
+        request_usage = streamed_response.usage()
 
         self._finish_handling(ctx, model_response, request_usage)
         assert self._result is not None  # this should be set by the previous line
 
     async def _make_request(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
-    ) -> HandleResponseNode[DepsT, NodeRunEndT]:
+    ) -> CallToolsNode[DepsT, NodeRunEndT]:
         if self._result is not None:
             return self._result
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
-        with _logfire.span('model request', run_step=ctx.state.run_step) as span:
-            model_response, request_usage = await ctx.deps.model.request(
-                ctx.state.message_history, model_settings, model_request_parameters
-            )
-            ctx.state.usage.incr(_usage.Usage(), requests=1)
-            span.set_attribute('response', model_response)
-            span.set_attribute('usage', request_usage)
+        model_response, request_usage = await ctx.deps.model.request(
+            ctx.state.message_history, model_settings, model_request_parameters
+        )
+        ctx.state.usage.incr(_usage.Usage(), requests=1)
 
         return self._finish_handling(ctx, model_response, request_usage)
 
@@ -344,7 +339,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
         response: _messages.ModelResponse,
         usage: _usage.Usage,
-    ) -> HandleResponseNode[DepsT, NodeRunEndT]:
+    ) -> CallToolsNode[DepsT, NodeRunEndT]:
         # Update usage
         ctx.state.usage.incr(usage, requests=0)
         if ctx.deps.usage_limits:
@@ -354,13 +349,13 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         ctx.state.message_history.append(response)
 
         # Set the `_result` attribute since we can't use `return` in an async iterator
-        self._result = HandleResponseNode(response)
+        self._result = CallToolsNode(response)
 
         return self._result
 
 
 @dataclasses.dataclass
-class HandleResponseNode(AgentNode[DepsT, NodeRunEndT]):
+class CallToolsNode(AgentNode[DepsT, NodeRunEndT]):
     """Process a model response, and decide whether to end the run or make a new request."""
 
     model_response: _messages.ModelResponse
@@ -454,8 +449,7 @@ class HandleResponseNode(AgentNode[DepsT, NodeRunEndT]):
         final_result: result.FinalResult[NodeRunEndT] | None = None
         parts: list[_messages.ModelRequestPart] = []
         if result_schema is not None:
-            if match := result_schema.find_tool(tool_calls):
-                call, result_tool = match
+            for call, result_tool in result_schema.find_tool(tool_calls):
                 try:
                     result_data = result_tool.validate(call)
                     result_data = await _validate_result(result_data, ctx, call)
@@ -465,12 +459,17 @@ class HandleResponseNode(AgentNode[DepsT, NodeRunEndT]):
                     ctx.state.increment_retries(ctx.deps.max_result_retries)
                     parts.append(e.tool_retry)
                 else:
-                    final_result = result.FinalResult(result_data, call.tool_name)
+                    final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
+                    break
 
         # Then build the other request parts based on end strategy
         tool_responses: list[_messages.ModelRequestPart] = self._tool_responses
         async for event in process_function_tools(
-            tool_calls, final_result and final_result.tool_name, ctx, tool_responses
+            tool_calls,
+            final_result and final_result.tool_name,
+            final_result and final_result.tool_call_id,
+            ctx,
+            tool_responses,
         ):
             yield event
 
@@ -496,7 +495,10 @@ class HandleResponseNode(AgentNode[DepsT, NodeRunEndT]):
             messages.append(_messages.ModelRequest(parts=tool_responses))
 
         run_span.set_attribute('usage', usage)
-        run_span.set_attribute('all_messages', messages)
+        run_span.set_attribute(
+            'all_messages_events',
+            [InstrumentedModel.event_to_dict(e) for e in InstrumentedModel.messages_to_otel_events(messages)],
+        )
 
         # End the run with self.data
         return End(final_result)
@@ -518,7 +520,7 @@ class HandleResponseNode(AgentNode[DepsT, NodeRunEndT]):
                 return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
             else:
                 # The following cast is safe because we know `str` is an allowed result type
-                return self._handle_final_result(ctx, result.FinalResult(result_data, tool_name=None), [])
+                return self._handle_final_result(ctx, result.FinalResult(result_data, None, None), [])
         else:
             ctx.state.increment_retries(ctx.deps.max_result_retries)
             return ModelRequestNode[DepsT, NodeRunEndT](
@@ -547,6 +549,7 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
 async def process_function_tools(
     tool_calls: list[_messages.ToolCallPart],
     result_tool_name: str | None,
+    result_tool_call_id: str | None,
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
     output_parts: list[_messages.ModelRequestPart],
 ) -> AsyncIterator[_messages.HandleResponseEvent]:
@@ -566,7 +569,11 @@ async def process_function_tools(
     calls_to_run: list[tuple[Tool[DepsT], _messages.ToolCallPart]] = []
     call_index_to_event_id: dict[int, str] = {}
     for call in tool_calls:
-        if call.tool_name == result_tool_name and not found_used_result_tool:
+        if (
+            call.tool_name == result_tool_name
+            and call.tool_call_id == result_tool_call_id
+            and not found_used_result_tool
+        ):
             found_used_result_tool = True
             output_parts.append(
                 _messages.ToolReturnPart(
@@ -593,9 +600,14 @@ async def process_function_tools(
             # if tool_name is in _result_schema, it means we found a result tool but an error occurred in
             # validation, we don't add another part here
             if result_tool_name is not None:
+                if found_used_result_tool:
+                    content = 'Result tool not used - a final result was already processed.'
+                else:
+                    # TODO: Include information about the validation failure, and/or merge this with the ModelRetry part
+                    content = 'Result tool not used - result failed validation.'
                 part = _messages.ToolReturnPart(
                     tool_name=call.tool_name,
-                    content='Result tool not used - a final result was already processed.',
+                    content=content,
                     tool_call_id=call.tool_call_id,
                 )
                 output_parts.append(part)
@@ -716,7 +728,7 @@ def build_agent_graph(
     nodes = (
         UserPromptNode[DepsT],
         ModelRequestNode[DepsT],
-        HandleResponseNode[DepsT],
+        CallToolsNode[DepsT],
     )
     graph = Graph[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[ResultT]](
         nodes=nodes,

pydantic_ai/_result.py

@@ -3,7 +3,7 @@ from __future__ import annotations as _annotations
 import inspect
 import sys
 import types
-from collections.abc import Awaitable, Iterable
+from collections.abc import Awaitable, Iterable, Iterator
 from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, Literal, Union, cast, get_args, get_origin
 
@@ -127,12 +127,12 @@ class ResultSchema(Generic[ResultDataT]):
     def find_tool(
         self,
         parts: Iterable[_messages.ModelResponsePart],
-    ) -> tuple[_messages.ToolCallPart, ResultTool[ResultDataT]] | None:
+    ) -> Iterator[tuple[_messages.ToolCallPart, ResultTool[ResultDataT]]]:
         """Find a tool that matches one of the calls."""
         for part in parts:
             if isinstance(part, _messages.ToolCallPart):
                 if result := self.tools.get(part.tool_name):
-                    return part, result
+                    yield part, result
 
     def tool_names(self) -> list[str]:
         """Return the names of the tools."""

pydantic_ai/_utils.py

@@ -48,6 +48,8 @@ def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema:
 
     if schema.get('type') == 'object':
         return schema
+    elif schema.get('$ref') is not None:
+        return schema.get('$defs', {}).get(schema['$ref'][8:])  # This removes the initial "#/$defs/".
     else:
         raise UserError('Schema must be an object')
 

pydantic_ai/agent.py

@@ -25,6 +25,7 @@ from . import (
     result,
     usage as _usage,
 )
+from .models.instrumented import InstrumentedModel
 from .result import FinalResult, ResultDataT, StreamedRunResult
 from .settings import ModelSettings, merge_model_settings
 from .tools import (
@@ -42,7 +43,7 @@ from .tools import (
 # Re-exporting like this improves auto-import behavior in PyCharm
 capture_run_messages = _agent_graph.capture_run_messages
 EndStrategy = _agent_graph.EndStrategy
-HandleResponseNode = _agent_graph.HandleResponseNode
+CallToolsNode = _agent_graph.CallToolsNode
 ModelRequestNode = _agent_graph.ModelRequestNode
 UserPromptNode = _agent_graph.UserPromptNode
 
@@ -52,7 +53,7 @@ __all__ = (
     'AgentRunResult',
     'capture_run_messages',
     'EndStrategy',
-    'HandleResponseNode',
+    'CallToolsNode',
     'ModelRequestNode',
     'UserPromptNode',
 )
@@ -294,7 +295,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
-        with self.iter(
+        async with self.iter(
             user_prompt=user_prompt,
             result_type=result_type,
             message_history=message_history,
@@ -310,8 +311,8 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         assert (final_result := agent_run.result) is not None, 'The graph run did not finish properly'
         return final_result
 
-    @contextmanager
-    def iter(
+    @asynccontextmanager
+    async def iter(
         self,
         user_prompt: str | Sequence[_messages.UserContent],
         *,
@@ -323,7 +324,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> Iterator[AgentRun[AgentDepsT, Any]]:
+    ) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
         """A contextmanager which can be used to iterate over the agent graph's nodes as they are executed.
 
         This method builds an internal agent graph (using system prompts, tools and result schemas) and then returns an
@@ -344,7 +345,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 
         async def main():
             nodes = []
-            with agent.iter('What is the capital of France?') as agent_run:
+            async with agent.iter('What is the capital of France?') as agent_run:
                 async for node in agent_run:
                     nodes.append(node)
             print(nodes)
@@ -362,7 +363,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                         kind='request',
                     )
                 ),
-                HandleResponseNode(
+                CallToolsNode(
                     model_response=ModelResponse(
                         parts=[TextPart(content='Paris', part_kind='text')],
                         model_name='gpt-4o',
@@ -370,7 +371,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                         kind='response',
                     )
                 ),
-                End(data=FinalResult(data='Paris', tool_name=None)),
+                End(data=FinalResult(data='Paris', tool_name=None, tool_call_id=None)),
             ]
             '''
             print(agent_run.result.data)
@@ -454,7 +455,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
             system_prompt_dynamic_functions=self._system_prompt_dynamic_functions,
         )
 
-        with graph.iter(
+        async with graph.iter(
             start_node,
             state=state,
             deps=graph_deps,
@@ -633,7 +634,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                 self._infer_name(frame.f_back)
 
         yielded = False
-        with self.iter(
+        async with self.iter(
             user_prompt,
             result_type=result_type,
             message_history=message_history,
@@ -661,11 +662,10 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                                     new_part = maybe_part_event.part
                                     if isinstance(new_part, _messages.TextPart):
                                         if _agent_graph.allow_text_result(result_schema):
-                                            return FinalResult(s, None)
-                                    elif isinstance(new_part, _messages.ToolCallPart):
-                                        if result_schema is not None and (match := result_schema.find_tool([new_part])):
-                                            call, _ = match
-                                            return FinalResult(s, call.tool_name)
+                                            return FinalResult(s, None, None)
+                                    elif isinstance(new_part, _messages.ToolCallPart) and result_schema:
+                                        for call, _ in result_schema.find_tool([new_part]):
+                                            return FinalResult(s, call.tool_name, call.tool_call_id)
                             return None
 
                         final_result_details = await stream_to_final(streamed_response)
@@ -692,6 +692,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                                 async for _event in _agent_graph.process_function_tools(
                                     tool_calls,
                                     final_result_details.tool_name,
+                                    final_result_details.tool_call_id,
                                     graph_ctx,
                                     parts,
                                 ):
@@ -1115,6 +1116,9 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             raise exceptions.UserError('`model` must be set either when creating the agent or when calling it.')
 
+        if not isinstance(model_, InstrumentedModel):
+            model_ = InstrumentedModel(model_)
+
         return model_
 
     def _get_deps(self: Agent[T, ResultDataT], deps: T) -> T:
@@ -1183,14 +1187,14 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         return isinstance(node, _agent_graph.ModelRequestNode)
 
     @staticmethod
-    def is_handle_response_node(
+    def is_call_tools_node(
         node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],
-    ) -> TypeGuard[_agent_graph.HandleResponseNode[T, S]]:
-        """Check if the node is a `HandleResponseNode`, narrowing the type if it is.
+    ) -> TypeGuard[_agent_graph.CallToolsNode[T, S]]:
+        """Check if the node is a `CallToolsNode`, narrowing the type if it is.
 
         This method preserves the generic parameters while narrowing the type, unlike a direct call to `isinstance`.
         """
-        return isinstance(node, _agent_graph.HandleResponseNode)
+        return isinstance(node, _agent_graph.CallToolsNode)
 
     @staticmethod
     def is_user_prompt_node(
@@ -1217,7 +1221,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
 class AgentRun(Generic[AgentDepsT, ResultDataT]):
     """A stateful, async-iterable run of an [`Agent`][pydantic_ai.agent.Agent].
 
-    You generally obtain an `AgentRun` instance by calling `with my_agent.iter(...) as agent_run:`.
+    You generally obtain an `AgentRun` instance by calling `async with my_agent.iter(...) as agent_run:`.
 
     Once you have an instance, you can use it to iterate through the run's nodes as they execute. When an
     [`End`][pydantic_graph.nodes.End] is reached, the run finishes and [`result`][pydantic_ai.agent.AgentRun.result]
@@ -1232,7 +1236,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
     async def main():
         nodes = []
         # Iterate through the run, recording each node along the way:
-        with agent.iter('What is the capital of France?') as agent_run:
+        async with agent.iter('What is the capital of France?') as agent_run:
             async for node in agent_run:
                 nodes.append(node)
         print(nodes)
@@ -1250,7 +1254,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
                     kind='request',
                 )
             ),
-            HandleResponseNode(
+            CallToolsNode(
                 model_response=ModelResponse(
                     parts=[TextPart(content='Paris', part_kind='text')],
                     model_name='gpt-4o',
@@ -1258,7 +1262,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
                     kind='response',
                 )
             ),
-            End(data=FinalResult(data='Paris', tool_name=None)),
+            End(data=FinalResult(data='Paris', tool_name=None, tool_call_id=None)),
         ]
         '''
         print(agent_run.result.data)
@@ -1346,7 +1350,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
         agent = Agent('openai:gpt-4o')
 
         async def main():
-            with agent.iter('What is the capital of France?') as agent_run:
+            async with agent.iter('What is the capital of France?') as agent_run:
                 next_node = agent_run.next_node  # start with the first node
                 nodes = [next_node]
                 while not isinstance(next_node, End):
@@ -1374,7 +1378,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
                             kind='request',
                         )
                     ),
-                    HandleResponseNode(
+                    CallToolsNode(
                         model_response=ModelResponse(
                             parts=[TextPart(content='Paris', part_kind='text')],
                             model_name='gpt-4o',
@@ -1382,7 +1386,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
                             kind='response',
                         )
                     ),
-                    End(data=FinalResult(data='Paris', tool_name=None)),
+                    End(data=FinalResult(data='Paris', tool_name=None, tool_call_id=None)),
                 ]
                 '''
                 print('Final result:', agent_run.result.data)

pydantic_ai/messages.py

@@ -8,6 +8,7 @@ from typing import Annotated, Any, Literal, Union, cast, overload
 
 import pydantic
 import pydantic_core
+from opentelemetry._events import Event
 from typing_extensions import TypeAlias
 
 from ._utils import now_utc as _now_utc
@@ -33,6 +34,9 @@ class SystemPromptPart:
     part_kind: Literal['system-prompt'] = 'system-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    def otel_event(self) -> Event:
+        return Event('gen_ai.system.message', body={'content': self.content, 'role': 'system'})
+
 
 @dataclass
 class AudioUrl:
@@ -138,6 +142,14 @@ class UserPromptPart:
     part_kind: Literal['user-prompt'] = 'user-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
+    def otel_event(self) -> Event:
+        if isinstance(self.content, str):
+            content = self.content
+        else:
+            # TODO figure out what to record for images and audio
+            content = [part if isinstance(part, str) else {'kind': part.kind} for part in self.content]
+        return Event('gen_ai.user.message', body={'content': content, 'role': 'user'})
+
 
 tool_return_ta: pydantic.TypeAdapter[Any] = pydantic.TypeAdapter(Any, config=pydantic.ConfigDict(defer_build=True))
 
@@ -176,6 +188,9 @@ class ToolReturnPart:
         else:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
+    def otel_event(self) -> Event:
+        return Event('gen_ai.tool.message', body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id})
+
 
 error_details_ta = pydantic.TypeAdapter(list[pydantic_core.ErrorDetails], config=pydantic.ConfigDict(defer_build=True))
 
@@ -224,6 +239,14 @@ class RetryPromptPart:
             description = f'{len(self.content)} validation errors: {json_errors.decode()}'
         return f'{description}\n\nFix the errors and try again.'
 
+    def otel_event(self) -> Event:
+        if self.tool_name is None:
+            return Event('gen_ai.user.message', body={'content': self.model_response(), 'role': 'user'})
+        else:
+            return Event(
+                'gen_ai.tool.message', body={'content': self.model_response(), 'role': 'tool', 'id': self.tool_call_id}
+            )
+
 
 ModelRequestPart = Annotated[
     Union[SystemPromptPart, UserPromptPart, ToolReturnPart, RetryPromptPart], pydantic.Discriminator('part_kind')
@@ -329,6 +352,36 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
+    def otel_events(self) -> list[Event]:
+        """Return OpenTelemetry events for the response."""
+        result: list[Event] = []
+
+        def new_event_body():
+            new_body: dict[str, Any] = {'role': 'assistant'}
+            ev = Event('gen_ai.assistant.message', body=new_body)
+            result.append(ev)
+            return new_body
+
+        body = new_event_body()
+        for part in self.parts:
+            if isinstance(part, ToolCallPart):
+                body.setdefault('tool_calls', []).append(
+                    {
+                        'id': part.tool_call_id,
+                        'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
+                        'function': {
+                            'name': part.tool_name,
+                            'arguments': part.args,
+                        },
+                    }
+                )
+            elif isinstance(part, TextPart):
+                if body.get('content'):
+                    body = new_event_body()
+                body['content'] = part.content
+
+        return result
+
 
 ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]
 """Any message sent to or returned by a model."""
@@ -539,6 +592,8 @@ class FinalResultEvent:
 
     tool_name: str | None
     """The name of the result tool that was called. `None` if the result is from text content and not from a tool."""
+    tool_call_id: str | None
+    """The tool call ID, if any, that this result is associated with."""
     event_kind: Literal['final_result'] = 'final_result'
     """Event type identifier, used as a discriminator."""
 

pydantic_ai/models/instrumented.py

@@ -1,28 +1,21 @@
 from __future__ import annotations
 
 import json
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
-from functools import partial
 from typing import Any, Callable, Literal
 
 import logfire_api
 from opentelemetry._events import Event, EventLogger, EventLoggerProvider, get_event_logger_provider
-from opentelemetry.trace import Tracer, TracerProvider, get_tracer_provider
+from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provider
 from opentelemetry.util.types import AttributeValue
+from pydantic import TypeAdapter
 
 from ..messages import (
     ModelMessage,
     ModelRequest,
-    ModelRequestPart,
     ModelResponse,
-    RetryPromptPart,
-    SystemPromptPart,
-    TextPart,
-    ToolCallPart,
-    ToolReturnPart,
-    UserPromptPart,
 )
 from ..settings import ModelSettings
 from ..usage import Usage
@@ -48,6 +41,8 @@ MODEL_SETTING_ATTRIBUTES: tuple[
     'frequency_penalty',
 )
 
+ANY_ADAPTER = TypeAdapter[Any](Any)
+
 
 @dataclass
 class InstrumentedModel(WrapperModel):
@@ -115,7 +110,7 @@ class InstrumentedModel(WrapperModel):
                     finish(response_stream.get(), response_stream.usage())
 
     @contextmanager
-    def _instrument(  # noqa: C901
+    def _instrument(
         self,
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
@@ -141,35 +136,24 @@ class InstrumentedModel(WrapperModel):
                 if isinstance(value := model_settings.get(key), (float, int)):
                     attributes[f'gen_ai.request.{key}'] = value
 
-        events_list = []
-        emit_event = partial(self._emit_event, system, events_list)
-
         with self.tracer.start_as_current_span(span_name, attributes=attributes) as span:
-            if span.is_recording():
-                for message in messages:
-                    if isinstance(message, ModelRequest):
-                        for part in message.parts:
-                            event_name, body = _request_part_body(part)
-                            if event_name:
-                                emit_event(event_name, body)
-                    elif isinstance(message, ModelResponse):
-                        for body in _response_bodies(message):
-                            emit_event('gen_ai.assistant.message', body)
 
             def finish(response: ModelResponse, usage: Usage):
                 if not span.is_recording():
                     return
 
-                for response_body in _response_bodies(response):
-                    if response_body:
-                        emit_event(
+                events = self.messages_to_otel_events(messages)
+                for event in self.messages_to_otel_events([response]):
+                    events.append(
+                        Event(
                             'gen_ai.choice',
-                            {
+                            body={
                                 # TODO finish_reason
                                 'index': 0,
-                                'message': response_body,
+                                'message': event.body,
                             },
                         )
+                    )
                 span.set_attributes(
                     {
                         # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
@@ -178,67 +162,56 @@ class InstrumentedModel(WrapperModel):
                         **usage.opentelemetry_attributes(),
                     }
                 )
-                if events_list:
-                    attr_name = 'events'
-                    span.set_attributes(
-                        {
-                            attr_name: json.dumps(events_list),
-                            'logfire.json_schema': json.dumps(
-                                {
-                                    'type': 'object',
-                                    'properties': {attr_name: {'type': 'array'}},
-                                }
-                            ),
-                        }
-                    )
+                self._emit_events(system, span, events)
 
             yield finish
 
-    def _emit_event(
-        self, system: str, events_list: list[dict[str, Any]], event_name: str, body: dict[str, Any]
-    ) -> None:
-        attributes = {'gen_ai.system': system}
+    def _emit_events(self, system: str, span: Span, events: list[Event]) -> None:
+        for event in events:
+            event.attributes = {'gen_ai.system': system, **(event.attributes or {})}
         if self.event_mode == 'logs':
-            self.event_logger.emit(Event(event_name, body=body, attributes=attributes))
-        else:
-            events_list.append({'event.name': event_name, **body, **attributes})
-
-
-def _request_part_body(part: ModelRequestPart) -> tuple[str, dict[str, Any]]:
-    if isinstance(part, SystemPromptPart):
-        return 'gen_ai.system.message', {'content': part.content, 'role': 'system'}
-    elif isinstance(part, UserPromptPart):
-        return 'gen_ai.user.message', {'content': part.content, 'role': 'user'}
-    elif isinstance(part, ToolReturnPart):
-        return 'gen_ai.tool.message', {'content': part.content, 'role': 'tool', 'id': part.tool_call_id}
-    elif isinstance(part, RetryPromptPart):
-        if part.tool_name is None:
-            return 'gen_ai.user.message', {'content': part.model_response(), 'role': 'user'}
+            for event in events:
+                self.event_logger.emit(event)
         else:
-            return 'gen_ai.tool.message', {'content': part.model_response(), 'role': 'tool', 'id': part.tool_call_id}
-    else:
-        return '', {}
-
-
-def _response_bodies(message: ModelResponse) -> list[dict[str, Any]]:
-    body: dict[str, Any] = {'role': 'assistant'}
-    result = [body]
-    for part in message.parts:
-        if isinstance(part, ToolCallPart):
-            body.setdefault('tool_calls', []).append(
+            attr_name = 'events'
+            span.set_attributes(
                 {
-                    'id': part.tool_call_id,
-                    'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
-                    'function': {
-                        'name': part.tool_name,
-                        'arguments': part.args,
-                    },
+                    attr_name: json.dumps([self.event_to_dict(event) for event in events]),
+                    'logfire.json_schema': json.dumps(
+                        {
+                            'type': 'object',
+                            'properties': {attr_name: {'type': 'array'}},
+                        }
+                    ),
                 }
             )
-        elif isinstance(part, TextPart):
-            if body.get('content'):
-                body = {'role': 'assistant'}
-                result.append(body)
-            body['content'] = part.content
 
-    return result
+    @staticmethod
+    def event_to_dict(event: Event) -> dict[str, Any]:
+        if not event.body:
+            body = {}
+        elif isinstance(event.body, Mapping):
+            body = event.body  # type: ignore
+        else:
+            body = {'body': event.body}
+        return {**body, **(event.attributes or {})}
+
+    @staticmethod
+    def messages_to_otel_events(messages: list[ModelMessage]) -> list[Event]:
+        result: list[Event] = []
+        for message in messages:
+            if isinstance(message, ModelRequest):
+                for part in message.parts:
+                    if hasattr(part, 'otel_event'):
+                        result.append(part.otel_event())
+            elif isinstance(message, ModelResponse):
+                result.extend(message.otel_events())
+        for event in result:
+            try:
+                event.body = ANY_ADAPTER.dump_python(event.body, mode='json')
+            except Exception:
+                try:
+                    event.body = str(event.body)
+                except Exception:
+                    event.body = 'Unable to serialize event body'
+        return result

pydantic_ai/result.py

@@ -145,12 +145,14 @@ class AgentStream(Generic[AgentDepsT, ResultDataT]):
                 if isinstance(e, _messages.PartStartEvent):
                     new_part = e.part
                     if isinstance(new_part, _messages.ToolCallPart):
-                        if result_schema is not None and (match := result_schema.find_tool([new_part])):
-                            call, _ = match
-                            return _messages.FinalResultEvent(tool_name=call.tool_name)
+                        if result_schema:
+                            for call, _ in result_schema.find_tool([new_part]):
+                                return _messages.FinalResultEvent(
+                                    tool_name=call.tool_name, tool_call_id=call.tool_call_id
+                                )
                     elif allow_text_result:
                         assert_type(e, _messages.PartStartEvent)
-                        return _messages.FinalResultEvent(tool_name=None)
+                        return _messages.FinalResultEvent(tool_name=None, tool_call_id=None)
 
             usage_checking_stream = _get_usage_checking_stream_response(
                 self._raw_stream_response, self._usage_limits, self.usage
@@ -472,6 +474,8 @@ class FinalResult(Generic[ResultDataT]):
     """The final result data."""
     tool_name: str | None
     """Name of the final result tool; `None` if the result came from unstructured text content."""
+    tool_call_id: str | None
+    """ID of the tool call that produced the final result; `None` if the result came from unstructured text content."""
 
 
 def _get_usage_checking_stream_response(

{pydantic_ai_slim-0.0.30.dist-info → pydantic_ai_slim-0.0.31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.30
+Version: 0.0.31
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,7 +29,8 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: logfire-api>=1.2.0
-Requires-Dist: pydantic-graph==0.0.30
+Requires-Dist: opentelemetry-api>=1.28.0
+Requires-Dist: pydantic-graph==0.0.31
 Requires-Dist: pydantic>=2.10
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.40.0; extra == 'anthropic'

{pydantic_ai_slim-0.0.30.dist-info → pydantic_ai_slim-0.0.31.dist-info}/RECORD

@@ -1,17 +1,17 @@
-pydantic_ai/__init__.py,sha256=Rmpjmorf8YY1PtlkXRRNN-J3ZoQDSh7chaibVGyqY0k,937
-pydantic_ai/_agent_graph.py,sha256=gvJQ17A2glk8p2w2TCSfHwvWNp0vla1sQb0EZXOZbxU,30284
+pydantic_ai/__init__.py,sha256=xrSDxkBwpUVInbPtTVhReEecStk-mWZMttAPUAQR0Ic,927
+pydantic_ai/_agent_graph.py,sha256=vvhV051rjVcPPRZ_TeL4pWwX-DptEzWgBBJnhybmIWg,30510
 pydantic_ai/_griffe.py,sha256=RYRKiLbgG97QxnazbAwlnc74XxevGHLQet-FGfq9qls,3960
 pydantic_ai/_parts_manager.py,sha256=ARfDQY1_5AIY5rNl_M2fAYHEFCe03ZxdhgjHf9qeIKw,11872
 pydantic_ai/_pydantic.py,sha256=dROz3Hmfdi0C2exq88FhefDRVo_8S3rtkXnoUHzsz0c,8753
-pydantic_ai/_result.py,sha256=tN1pVulf_EM4bkBvpNUWPnUXezLY-sBrJEVCFdy2nLU,10264
+pydantic_ai/_result.py,sha256=mqj3YrUzr5OT00h0KfGJglwQZ6_7nV7355Pvucd08ak,10276
 pydantic_ai/_system_prompt.py,sha256=602c2jyle2R_SesOrITBDETZqsLk4BZ8Cbo8yEhmx04,1120
-pydantic_ai/_utils.py,sha256=w9BYSfFZiaX757fRtMRclOL1uYzyQnxV_lxqbU2WTPs,9435
-pydantic_ai/agent.py,sha256=FeKELTSFKDkt6-UlmkezKnQTdnx1in6VckivqsfzfA4,65382
+pydantic_ai/_utils.py,sha256=nx4Suswk2qjLvzphx8uQntKzFi-IzvhX_H1L7t_kJlQ,9579
+pydantic_ai/agent.py,sha256=jHQ99M-kwUrUSWHPjBDmWG2AepbDS9H3YUE1NugaWGg,65625
 pydantic_ai/exceptions.py,sha256=1ujJeB3jDDQ-pH5ydBYrgStvR35-GlEW0bYGTGEr4ME,3127
 pydantic_ai/format_as_xml.py,sha256=QE7eMlg5-YUMw1_2kcI3h0uKYPZZyGkgXFDtfZTMeeI,4480
-pydantic_ai/messages.py,sha256=k8sX-V1cTeqXh1u6oJbqExZPYt3E7F3UCIudxvjKRO8,21486
+pydantic_ai/messages.py,sha256=Yny2hIuExXfw9fvHDSPgbvfN91IOdcLaDEAmaCAoTBs,23751
 pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai/result.py,sha256=Df_tPeqCQnLa0i0vVA-BGCJDx37ebD_3ojAmHnXE2yU,22767
+pydantic_ai/result.py,sha256=Q--JTwDfPeJw1_Mk5EhI7R9V7GusG-oAx1m9pDH50zQ,23014
 pydantic_ai/settings.py,sha256=ntuWnke9UA18aByDxk9OIhN0tAgOaPdqCEkRf-wlp8Y,3059
 pydantic_ai/tools.py,sha256=IPZuZJCSQUppz1uyLVwpfFLGoMirB8YtKWXIDQGR444,13414
 pydantic_ai/usage.py,sha256=VmpU_o_RjFI65J81G1wfCwDIAYBclMjeWfLtslntFOw,5406
@@ -25,12 +25,12 @@ pydantic_ai/models/fallback.py,sha256=smHwNIpxu19JsgYYjY0nmzl3yox7yQRJ0Ir08zdhnk
 pydantic_ai/models/function.py,sha256=THIwVJ8qI3efYLNlYXlYze_J8hc7MHB-NMb3kpknq0g,11373
 pydantic_ai/models/gemini.py,sha256=2hDTMIMf899dp-MS0tLT7m1GkXsL9KIRMBklGM0VLB4,34223
 pydantic_ai/models/groq.py,sha256=Z4sZJDu5Yxa2tZiAPp9EjSVMz4uwLhS3fW7kFSc09gI,16406
-pydantic_ai/models/instrumented.py,sha256=xUZEn2VG8hP3hny0L5kZgXC5UnFdlUJ0DgXOxFmYhEo,9654
+pydantic_ai/models/instrumented.py,sha256=npufEZJrR9m0_ZQB1inWFcuK3Nu5_2GdY1YtTYaIj3s,8366
 pydantic_ai/models/mistral.py,sha256=ZJ4xPcL9wJIQ5io34yP2fPyXy8GZrSvsW4itZiKPYFw,27448
 pydantic_ai/models/openai.py,sha256=koIcK_pDHmV-JFq_-VIzU-edAqGKOOzkSk5QSYWvfoc,20156
 pydantic_ai/models/test.py,sha256=Ux20cmuJFkhvI9L1N7ItHNFcd-j284TBEsrM53eWRag,16873
 pydantic_ai/models/vertexai.py,sha256=9Kp_1KMBlbP8_HRJTuFnrkkFmlJ7yFhADQYjxOgIh9Y,9523
 pydantic_ai/models/wrapper.py,sha256=Zr3fgiUBpt2N9gXds6iSwaMEtEsFKr9WwhpHjSoHa7o,1410
-pydantic_ai_slim-0.0.30.dist-info/METADATA,sha256=JDT77S9uw0w87WpAbXqK_c65849A7PeF1_dhJRGamiM,3062
-pydantic_ai_slim-0.0.30.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydantic_ai_slim-0.0.30.dist-info/RECORD,,
+pydantic_ai_slim-0.0.31.dist-info/METADATA,sha256=dgkUKEU7r9OqgIkt3enzpISWt73KVAYL8gC2APlnpWg,3103
+pydantic_ai_slim-0.0.31.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydantic_ai_slim-0.0.31.dist-info/RECORD,,