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

pydantic_ai/agent.py

@@ -9,9 +9,9 @@ from types import FrameType
 from typing import Any, Callable, Generic, cast, final, overload
 
 import logfire_api
-from typing_extensions import TypeVar, deprecated
+from typing_extensions import TypeGuard, TypeVar, deprecated
 
-from pydantic_graph import BaseNode, End, Graph, GraphRun, GraphRunContext
+from pydantic_graph import End, Graph, GraphRun, GraphRunContext
 from pydantic_graph._utils import get_event_loop
 
 from . import (
@@ -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,18 +43,17 @@ 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
 
-
 __all__ = (
     'Agent',
     'AgentRun',
     'AgentRunResult',
     'capture_run_messages',
     'EndStrategy',
-    'HandleResponseNode',
+    'CallToolsNode',
     'ModelRequestNode',
     'UserPromptNode',
 )
@@ -71,6 +71,7 @@ else:
     logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
 T = TypeVar('T')
+S = TypeVar('S')
 NoneType = type(None)
 RunResultDataT = TypeVar('RunResultDataT')
 """Type variable for the result data of a run where `result_type` was customized on the run call."""
@@ -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,
@@ -646,10 +647,9 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         ) as agent_run:
             first_node = agent_run.next_node  # start with the first node
             assert isinstance(first_node, _agent_graph.UserPromptNode)  # the first node should be a user prompt node
-            node: BaseNode[Any, Any, Any] = cast(BaseNode[Any, Any, Any], first_node)
+            node = first_node
             while True:
-                if isinstance(node, _agent_graph.ModelRequestNode):
-                    node = cast(_agent_graph.ModelRequestNode[AgentDepsT, Any], node)
+                if self.is_model_request_node(node):
                     graph_ctx = agent_run.ctx
                     async with node._stream(graph_ctx) as streamed_response:  # pyright: ignore[reportPrivateUsage]
 
@@ -662,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)
@@ -693,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,
                                 ):
@@ -717,9 +717,9 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
                             )
                             break
                 next_node = await agent_run.next(node)
-                if not isinstance(next_node, BaseNode):
+                if not isinstance(next_node, _agent_graph.AgentNode):
                     raise exceptions.AgentRunError('Should have produced a StreamedRunResult before getting here')
-                node = cast(BaseNode[Any, Any, Any], next_node)
+                node = cast(_agent_graph.AgentNode[Any, Any], next_node)
 
         if not yielded:
             raise exceptions.AgentRunError('Agent run finished without producing a final result')
@@ -1116,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:
@@ -1173,12 +1176,52 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         else:
             return self._result_schema  # pyright: ignore[reportReturnType]
 
+    @staticmethod
+    def is_model_request_node(
+        node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],
+    ) -> TypeGuard[_agent_graph.ModelRequestNode[T, S]]:
+        """Check if the node is a `ModelRequestNode`, 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.ModelRequestNode)
+
+    @staticmethod
+    def is_call_tools_node(
+        node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],
+    ) -> 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.CallToolsNode)
+
+    @staticmethod
+    def is_user_prompt_node(
+        node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],
+    ) -> TypeGuard[_agent_graph.UserPromptNode[T, S]]:
+        """Check if the node is a `UserPromptNode`, 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.UserPromptNode)
+
+    @staticmethod
+    def is_end_node(
+        node: _agent_graph.AgentNode[T, S] | End[result.FinalResult[S]],
+    ) -> TypeGuard[End[result.FinalResult[S]]]:
+        """Check if the node is a `End`, 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, End)
+
 
 @dataclasses.dataclass(repr=False)
 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]
@@ -1193,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)
@@ -1211,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',
@@ -1219,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)
@@ -1244,15 +1287,17 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
     @property
     def next_node(
         self,
-    ) -> (
-        BaseNode[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any], FinalResult[ResultDataT]]
-        | End[FinalResult[ResultDataT]]
-    ):
+    ) -> _agent_graph.AgentNode[AgentDepsT, ResultDataT] | End[FinalResult[ResultDataT]]:
         """The next node that will be run in the agent graph.
 
         This is the next node that will be used during async iteration, or if a node is not passed to `self.next(...)`.
         """
-        return self._graph_run.next_node
+        next_node = self._graph_run.next_node
+        if isinstance(next_node, End):
+            return next_node
+        if _agent_graph.is_agent_node(next_node):
+            return next_node
+        raise exceptions.AgentRunError(f'Unexpected node type: {type(next_node)}')  # pragma: no cover
 
     @property
     def result(self) -> AgentRunResult[ResultDataT] | None:
@@ -1273,45 +1318,24 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
 
     def __aiter__(
         self,
-    ) -> AsyncIterator[
-        BaseNode[
-            _agent_graph.GraphAgentState,
-            _agent_graph.GraphAgentDeps[AgentDepsT, Any],
-            FinalResult[ResultDataT],
-        ]
-        | End[FinalResult[ResultDataT]]
-    ]:
+    ) -> AsyncIterator[_agent_graph.AgentNode[AgentDepsT, ResultDataT] | End[FinalResult[ResultDataT]]]:
         """Provide async-iteration over the nodes in the agent run."""
         return self
 
     async def __anext__(
         self,
-    ) -> (
-        BaseNode[
-            _agent_graph.GraphAgentState,
-            _agent_graph.GraphAgentDeps[AgentDepsT, Any],
-            FinalResult[ResultDataT],
-        ]
-        | End[FinalResult[ResultDataT]]
-    ):
+    ) -> _agent_graph.AgentNode[AgentDepsT, ResultDataT] | End[FinalResult[ResultDataT]]:
         """Advance to the next node automatically based on the last returned node."""
-        return await self._graph_run.__anext__()
+        next_node = await self._graph_run.__anext__()
+        if _agent_graph.is_agent_node(next_node):
+            return next_node
+        assert isinstance(next_node, End), f'Unexpected node type: {type(next_node)}'
+        return next_node
 
     async def next(
         self,
-        node: BaseNode[
-            _agent_graph.GraphAgentState,
-            _agent_graph.GraphAgentDeps[AgentDepsT, Any],
-            FinalResult[ResultDataT],
-        ],
-    ) -> (
-        BaseNode[
-            _agent_graph.GraphAgentState,
-            _agent_graph.GraphAgentDeps[AgentDepsT, Any],
-            FinalResult[ResultDataT],
-        ]
-        | End[FinalResult[ResultDataT]]
-    ):
+        node: _agent_graph.AgentNode[AgentDepsT, ResultDataT],
+    ) -> _agent_graph.AgentNode[AgentDepsT, ResultDataT] | End[FinalResult[ResultDataT]]:
         """Manually drive the agent run by passing in the node you want to run next.
 
         This lets you inspect or mutate the node before continuing execution, or skip certain nodes
@@ -1326,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):
@@ -1354,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',
@@ -1362,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)
@@ -1378,7 +1402,11 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
         """
         # Note: It might be nice to expose a synchronous interface for iteration, but we shouldn't do it
         # on this class, or else IDEs won't warn you if you accidentally use `for` instead of `async for` to iterate.
-        return await self._graph_run.next(node)
+        next_node = await self._graph_run.next(node)
+        if _agent_graph.is_agent_node(next_node):
+            return next_node
+        assert isinstance(next_node, End), f'Unexpected node type: {type(next_node)}'
+        return next_node
 
     def usage(self) -> _usage.Usage:
         """Get usage statistics for the run so far, including token usage, model requests, and so on."""

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."""
@@ -533,9 +586,26 @@ class PartDeltaEvent:
     """Event type identifier, used as a discriminator."""
 
 
+@dataclass
+class FinalResultEvent:
+    """An event indicating the response to the current model request matches the result schema."""
+
+    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."""
+
+
 ModelResponseStreamEvent = Annotated[Union[PartStartEvent, PartDeltaEvent], pydantic.Discriminator('event_kind')]
 """An event in the model response stream, either starting a new part or applying a delta to an existing one."""
 
+AgentStreamEvent = Annotated[
+    Union[PartStartEvent, PartDeltaEvent, FinalResultEvent], pydantic.Discriminator('event_kind')
+]
+"""An event in the agent stream."""
+
 
 @dataclass
 class FunctionToolCallEvent:
@@ -558,7 +628,7 @@ class FunctionToolResultEvent:
 
     result: ToolReturnPart | RetryPromptPart
     """The result of the call to the function tool."""
-    call_id: str
+    tool_call_id: str
     """An ID used to match the result to its original call."""
     event_kind: Literal['function_tool_result'] = 'function_tool_result'
     """Event type identifier, used as a discriminator."""

pydantic_ai/models/__init__.py

@@ -84,6 +84,8 @@ KnownModelName = Literal[
     'gpt-4-turbo-2024-04-09',
     'gpt-4-turbo-preview',
     'gpt-4-vision-preview',
+    'gpt-4.5-preview',
+    'gpt-4.5-preview-2025-02-27',
     'gpt-4o',
     'gpt-4o-2024-05-13',
     'gpt-4o-2024-08-06',
@@ -138,6 +140,8 @@ KnownModelName = Literal[
     'openai:gpt-4-turbo-2024-04-09',
     'openai:gpt-4-turbo-preview',
     'openai:gpt-4-vision-preview',
+    'openai:gpt-4.5-preview',
+    'openai:gpt-4.5-preview-2025-02-27',
     'openai:gpt-4o',
     'openai:gpt-4o-2024-05-13',
     'openai:gpt-4o-2024-08-06',

pydantic_ai/models/function.py

@@ -177,6 +177,8 @@ class DeltaToolCall:
     """Incremental change to the name of the tool."""
     json_args: str | None = None
     """Incremental change to the arguments as JSON"""
+    tool_call_id: str | None = None
+    """Incremental change to the tool call ID."""
 
 
 DeltaToolCalls: TypeAlias = dict[int, DeltaToolCall]
@@ -224,7 +226,7 @@ class FunctionStreamedResponse(StreamedResponse):
                         vendor_part_id=dtc_index,
                         tool_name=delta_tool_call.name,
                         args=delta_tool_call.json_args,
-                        tool_call_id=None,
+                        tool_call_id=delta_tool_call.tool_call_id,
                     )
                     if maybe_event is not None:
                         yield maybe_event
@@ -280,7 +282,16 @@ def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
         return 0
     if isinstance(content, str):
         return len(re.split(r'[\s",.:]+', content.strip()))
-    # TODO(Marcelo): We need to study how we can estimate the tokens for these types of content.
     else:  # pragma: no cover
-        assert isinstance(content, (AudioUrl, ImageUrl, BinaryContent))
-        return 0
+        tokens = 0
+        for part in content:
+            if isinstance(part, str):
+                tokens += len(re.split(r'[\s",.:]+', part.strip()))
+            # TODO(Marcelo): We need to study how we can estimate the tokens for these types of content.
+            if isinstance(part, (AudioUrl, ImageUrl)):
+                tokens += 0
+            elif isinstance(part, BinaryContent):
+                tokens += len(part.data)
+            else:
+                tokens += 0
+        return tokens