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

pydantic_ai/_agent_graph.py

@@ -2,7 +2,6 @@ from __future__ import annotations as _annotations
 
 import asyncio
 import dataclasses
-from abc import ABC
 from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
@@ -10,7 +9,7 @@ from dataclasses import field
 from typing import Any, Generic, Literal, Union, cast
 
 import logfire_api
-from typing_extensions import TypeVar, assert_never
+from typing_extensions import TypeGuard, TypeVar, assert_never
 
 from pydantic_graph import BaseNode, Graph, GraphRunContext
 from pydantic_graph.nodes import End, NodeRunEndT
@@ -55,6 +54,7 @@ else:
     logfire._internal.stack_info.NON_USER_CODE_PREFIXES += (str(Path(__file__).parent.absolute()),)
 
 T = TypeVar('T')
+S = TypeVar('S')
 NoneType = type(None)
 EndStrategy = Literal['early', 'exhaustive']
 """The strategy for handling multiple tool calls when a final result is found.
@@ -107,8 +107,31 @@ class GraphAgentDeps(Generic[DepsT, ResultDataT]):
     run_span: logfire_api.LogfireSpan
 
 
+class AgentNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
+    """The base class for all agent nodes.
+
+    Using subclass of `BaseNode` for all nodes reduces the amount of boilerplate of generics everywhere
+    """
+
+
+def is_agent_node(
+    node: BaseNode[GraphAgentState, GraphAgentDeps[T, Any], result.FinalResult[S]] | End[result.FinalResult[S]],
+) -> TypeGuard[AgentNode[T, S]]:
+    """Check if the provided node is an instance of `AgentNode`.
+
+    Usage:
+
+        if is_agent_node(node):
+            # `node` is an AgentNode
+            ...
+
+    This method preserves the generic parameters on the narrowed type, unlike `isinstance(node, AgentNode)`.
+    """
+    return isinstance(node, AgentNode)
+
+
 @dataclasses.dataclass
-class UserPromptNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]], ABC):
+class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
     user_prompt: str | Sequence[_messages.UserContent]
 
     system_prompts: tuple[str, ...]
@@ -215,7 +238,7 @@ async def _prepare_request_parameters(
 
 
 @dataclasses.dataclass
-class ModelRequestNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
+class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
     """Make a request to the model using the last message in state.message_history."""
 
     request: _messages.ModelRequest
@@ -236,12 +259,30 @@ class ModelRequestNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], res
 
         return await self._make_request(ctx)
 
+    @asynccontextmanager
+    async def stream(
+        self,
+        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, T]],
+    ) -> AsyncIterator[result.AgentStream[DepsT, T]]:
+        async with self._stream(ctx) as streamed_response:
+            agent_stream = result.AgentStream[DepsT, T](
+                streamed_response,
+                ctx.deps.result_schema,
+                ctx.deps.result_validators,
+                build_run_context(ctx),
+                ctx.deps.usage_limits,
+            )
+            yield agent_stream
+            # 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 agent_stream:
+                pass
+
     @asynccontextmanager
     async def _stream(
         self,
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, T]],
     ) -> AsyncIterator[models.StreamedResponse]:
-        # TODO: Consider changing this to return something more similar to a `StreamedRunResult`, then make it public
         assert not self._did_stream, 'stream() should only be called once per node'
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
@@ -319,7 +360,7 @@ class ModelRequestNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], res
 
 
 @dataclasses.dataclass
-class HandleResponseNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
+class HandleResponseNode(AgentNode[DepsT, NodeRunEndT]):
     """Process a model response, and decide whether to end the run or make a new request."""
 
     model_response: _messages.ModelResponse
@@ -575,7 +616,7 @@ async def process_function_tools(
             for task in done:
                 index = tasks.index(task)
                 result = task.result()
-                yield _messages.FunctionToolResultEvent(result, call_id=call_index_to_event_id[index])
+                yield _messages.FunctionToolResultEvent(result, tool_call_id=call_index_to_event_id[index])
                 if isinstance(result, (_messages.ToolReturnPart, _messages.RetryPromptPart)):
                     results_by_index[index] = result
                 else:

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 (
@@ -46,7 +46,6 @@ HandleResponseNode = _agent_graph.HandleResponseNode
 ModelRequestNode = _agent_graph.ModelRequestNode
 UserPromptNode = _agent_graph.UserPromptNode
 
-
 __all__ = (
     'Agent',
     'AgentRun',
@@ -71,6 +70,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."""
@@ -646,10 +646,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]
 
@@ -717,9 +716,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')
@@ -1173,6 +1172,46 @@ 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_handle_response_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.
+
+        This method preserves the generic parameters while narrowing the type, unlike a direct call to `isinstance`.
+        """
+        return isinstance(node, _agent_graph.HandleResponseNode)
+
+    @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]):
@@ -1244,15 +1283,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 +1314,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
@@ -1378,7 +1398,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

@@ -533,9 +533,24 @@ 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."""
+    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 +573,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

pydantic_ai/models/instrumented.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 from collections.abc import AsyncIterator, Iterator
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
@@ -9,6 +10,7 @@ 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.util.types import AttributeValue
 
 from ..messages import (
     ModelMessage,
@@ -46,40 +48,42 @@ MODEL_SETTING_ATTRIBUTES: tuple[
     'frequency_penalty',
 )
 
-NOT_GIVEN = object()
-
 
 @dataclass
 class InstrumentedModel(WrapperModel):
-    """Model which is instrumented with logfire."""
+    """Model which is instrumented with OpenTelemetry."""
 
     tracer: Tracer = field(repr=False)
     event_logger: EventLogger = field(repr=False)
+    event_mode: Literal['attributes', 'logs'] = 'attributes'
 
     def __init__(
         self,
         wrapped: Model | KnownModelName,
         tracer_provider: TracerProvider | None = None,
         event_logger_provider: EventLoggerProvider | None = None,
+        event_mode: Literal['attributes', 'logs'] = 'attributes',
     ):
         super().__init__(wrapped)
         tracer_provider = tracer_provider or get_tracer_provider()
         event_logger_provider = event_logger_provider or get_event_logger_provider()
         self.tracer = tracer_provider.get_tracer('pydantic-ai')
         self.event_logger = event_logger_provider.get_event_logger('pydantic-ai')
+        self.event_mode = event_mode
 
     @classmethod
     def from_logfire(
         cls,
         wrapped: Model | KnownModelName,
         logfire_instance: logfire_api.Logfire = logfire_api.DEFAULT_LOGFIRE_INSTANCE,
+        event_mode: Literal['attributes', 'logs'] = 'attributes',
     ) -> InstrumentedModel:
         if hasattr(logfire_instance.config, 'get_event_logger_provider'):
             event_provider = logfire_instance.config.get_event_logger_provider()
         else:
             event_provider = None
         tracer_provider = logfire_instance.config.get_tracer_provider()
-        return cls(wrapped, tracer_provider, event_provider)
+        return cls(wrapped, tracer_provider, event_provider, event_mode)
 
     async def request(
         self,
@@ -111,7 +115,7 @@ class InstrumentedModel(WrapperModel):
                     finish(response_stream.get(), response_stream.usage())
 
     @contextmanager
-    def _instrument(
+    def _instrument(  # noqa: C901
         self,
         messages: list[ModelMessage],
         model_settings: ModelSettings | None,
@@ -126,7 +130,7 @@ class InstrumentedModel(WrapperModel):
         #  - server.port: to parse from the base_url
         #  - error.type: unclear if we should do something here or just always rely on span exceptions
         #  - gen_ai.request.stop_sequences/top_k: model_settings doesn't include these
-        attributes: dict[str, Any] = {
+        attributes: dict[str, AttributeValue] = {
             'gen_ai.operation.name': operation,
             'gen_ai.system': system,
             'gen_ai.request.model': model_name,
@@ -134,10 +138,11 @@ class InstrumentedModel(WrapperModel):
 
         if model_settings:
             for key in MODEL_SETTING_ATTRIBUTES:
-                if (value := model_settings.get(key, NOT_GIVEN)) is not NOT_GIVEN:
+                if isinstance(value := model_settings.get(key), (float, int)):
                     attributes[f'gen_ai.request.{key}'] = value
 
-        emit_event = partial(self._emit_event, system)
+        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():
@@ -167,22 +172,36 @@ class InstrumentedModel(WrapperModel):
                         )
                 span.set_attributes(
                     {
-                        k: v
-                        for k, v in {
-                            # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
-                            #  https://github.com/pydantic/pydantic-ai/issues/886
-                            'gen_ai.response.model': response.model_name or model_name,
-                            'gen_ai.usage.input_tokens': usage.request_tokens,
-                            'gen_ai.usage.output_tokens': usage.response_tokens,
-                        }.items()
-                        if v is not None
+                        # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
+                        #  https://github.com/pydantic/pydantic-ai/issues/886
+                        'gen_ai.response.model': response.model_name or model_name,
+                        **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'}},
+                                }
+                            ),
+                        }
+                    )
 
             yield finish
 
-    def _emit_event(self, system: str, event_name: str, body: dict[str, Any]) -> None:
-        self.event_logger.emit(Event(event_name, body=body, attributes={'gen_ai.system': system}))
+    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}
+        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]]:

pydantic_ai/result.py

@@ -7,9 +7,10 @@ from datetime import datetime
 from typing import Generic, Union, cast
 
 import logfire_api
-from typing_extensions import TypeVar
+from typing_extensions import TypeVar, assert_type
 
 from . import _result, _utils, exceptions, messages as _messages, models
+from .messages import AgentStreamEvent, FinalResultEvent
 from .tools import AgentDepsT, RunContext
 from .usage import Usage, UsageLimits
 
@@ -51,6 +52,125 @@ Usage `ResultValidatorFunc[AgentDepsT, T]`.
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai')
 
 
+@dataclass
+class AgentStream(Generic[AgentDepsT, ResultDataT]):
+    _raw_stream_response: models.StreamedResponse
+    _result_schema: _result.ResultSchema[ResultDataT] | None
+    _result_validators: list[_result.ResultValidator[AgentDepsT, ResultDataT]]
+    _run_ctx: RunContext[AgentDepsT]
+    _usage_limits: UsageLimits | None
+
+    _agent_stream_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
+    _final_result_event: FinalResultEvent | None = field(default=None, init=False)
+    _initial_run_ctx_usage: Usage = field(init=False)
+
+    def __post_init__(self):
+        self._initial_run_ctx_usage = copy(self._run_ctx.usage)
+
+    async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[ResultDataT]:
+        """Asynchronously stream the (validated) agent outputs."""
+        async for response in self.stream_responses(debounce_by=debounce_by):
+            if self._final_result_event is not None:
+                yield await self._validate_response(response, self._final_result_event.tool_name, allow_partial=True)
+        if self._final_result_event is not None:
+            yield await self._validate_response(
+                self._raw_stream_response.get(), self._final_result_event.tool_name, allow_partial=False
+            )
+
+    async def stream_responses(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
+        """Asynchronously stream the (unvalidated) model responses for the agent."""
+        # if the message currently has any parts with content, yield before streaming
+        msg = self._raw_stream_response.get()
+        for part in msg.parts:
+            if part.has_content():
+                yield msg
+                break
+
+        async with _utils.group_by_temporal(self, debounce_by) as group_iter:
+            async for _items in group_iter:
+                yield self._raw_stream_response.get()  # current state of the response
+
+    def usage(self) -> Usage:
+        """Return the usage of the whole run.
+
+        !!! note
+            This won't return the full usage until the stream is finished.
+        """
+        return self._initial_run_ctx_usage + self._raw_stream_response.usage()
+
+    async def _validate_response(
+        self, message: _messages.ModelResponse, result_tool_name: str | None, *, allow_partial: bool = False
+    ) -> ResultDataT:
+        """Validate a structured result message."""
+        if self._result_schema is not None and result_tool_name is not None:
+            match = self._result_schema.find_named_tool(message.parts, result_tool_name)
+            if match is None:
+                raise exceptions.UnexpectedModelBehavior(
+                    f'Invalid response, unable to find tool: {self._result_schema.tool_names()}'
+                )
+
+            call, result_tool = match
+            result_data = result_tool.validate(call, allow_partial=allow_partial, wrap_validation_errors=False)
+
+            for validator in self._result_validators:
+                result_data = await validator.validate(result_data, call, self._run_ctx)
+            return result_data
+        else:
+            text = '\n\n'.join(x.content for x in message.parts if isinstance(x, _messages.TextPart))
+            for validator in self._result_validators:
+                text = await validator.validate(
+                    text,
+                    None,
+                    self._run_ctx,
+                )
+            # Since there is no result tool, we can assume that str is compatible with ResultDataT
+            return cast(ResultDataT, text)
+
+    def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
+        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s.
+
+        This proxies the _raw_stream_response and sends all events to the agent stream, while also checking for matches
+        on the result schema and emitting a [`FinalResultEvent`][pydantic_ai.messages.FinalResultEvent] if/when the
+        first match is found.
+        """
+        if self._agent_stream_iterator is not None:
+            return self._agent_stream_iterator
+
+        async def aiter():
+            result_schema = self._result_schema
+            allow_text_result = result_schema is None or result_schema.allow_text_result
+
+            def _get_final_result_event(e: _messages.ModelResponseStreamEvent) -> _messages.FinalResultEvent | None:
+                """Return an appropriate FinalResultEvent if `e` corresponds to a part that will produce a final result."""
+                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)
+                    elif allow_text_result:
+                        assert_type(e, _messages.PartStartEvent)
+                        return _messages.FinalResultEvent(tool_name=None)
+
+            usage_checking_stream = _get_usage_checking_stream_response(
+                self._raw_stream_response, self._usage_limits, self.usage
+            )
+            async for event in usage_checking_stream:
+                yield event
+                if (final_result_event := _get_final_result_event(event)) is not None:
+                    self._final_result_event = final_result_event
+                    yield final_result_event
+                    break
+
+            # If we broke out of the above loop, we need to yield the rest of the events
+            # If we didn't, this will just be a no-op
+            async for event in usage_checking_stream:
+                yield event
+
+        self._agent_stream_iterator = aiter()
+        return self._agent_stream_iterator
+
+
 @dataclass
 class StreamedRunResult(Generic[AgentDepsT, ResultDataT]):
     """Result of a streamed run that returns structured data via a tool call."""

pydantic_ai/usage.py

@@ -56,6 +56,16 @@ class Usage:
         new_usage.incr(other)
         return new_usage
 
+    def opentelemetry_attributes(self) -> dict[str, int]:
+        """Get the token limits as OpenTelemetry attributes."""
+        result = {
+            'gen_ai.usage.input_tokens': self.request_tokens,
+            'gen_ai.usage.output_tokens': self.response_tokens,
+        }
+        for key, value in (self.details or {}).items():
+            result[f'gen_ai.usage.details.{key}'] = value
+        return {k: v for k, v in result.items() if v is not None}
+
 
 @dataclass
 class UsageLimits:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.29
+Version: 0.0.30
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT
@@ -29,7 +29,7 @@ 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.29
+Requires-Dist: pydantic-graph==0.0.30
 Requires-Dist: pydantic>=2.10
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.40.0; extra == 'anthropic'
@@ -44,7 +44,7 @@ Requires-Dist: logfire>=2.3; extra == 'logfire'
 Provides-Extra: mistral
 Requires-Dist: mistralai>=1.2.5; extra == 'mistral'
 Provides-Extra: openai
-Requires-Dist: openai>=1.61.0; extra == 'openai'
+Requires-Dist: openai>=1.65.1; extra == 'openai'
 Provides-Extra: tavily
 Requires-Dist: tavily-python>=0.5.0; extra == 'tavily'
 Provides-Extra: vertexai

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

@@ -1,36 +1,36 @@
 pydantic_ai/__init__.py,sha256=Rmpjmorf8YY1PtlkXRRNN-J3ZoQDSh7chaibVGyqY0k,937
-pydantic_ai/_agent_graph.py,sha256=lNKTtUyVY14M0WODP5K1NUaE9zJA716-9rZutapSg8A,29042
+pydantic_ai/_agent_graph.py,sha256=gvJQ17A2glk8p2w2TCSfHwvWNp0vla1sQb0EZXOZbxU,30284
 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/_system_prompt.py,sha256=602c2jyle2R_SesOrITBDETZqsLk4BZ8Cbo8yEhmx04,1120
 pydantic_ai/_utils.py,sha256=w9BYSfFZiaX757fRtMRclOL1uYzyQnxV_lxqbU2WTPs,9435
-pydantic_ai/agent.py,sha256=wCXvGwPykn-cYf1_4bR4XqYofqwI6X2lj4L-0-MbMp4,63685
+pydantic_ai/agent.py,sha256=FeKELTSFKDkt6-UlmkezKnQTdnx1in6VckivqsfzfA4,65382
 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=U-RgeRsMR-Ew6IoeBDrnQVONX9AwxyVd0aTnAxEA7EM,20918
+pydantic_ai/messages.py,sha256=k8sX-V1cTeqXh1u6oJbqExZPYt3E7F3UCIudxvjKRO8,21486
 pydantic_ai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pydantic_ai/result.py,sha256=Rqbog6efO1l_bFJSuAd-_ZZLoQa_rz4motOGeR_5N3I,16803
+pydantic_ai/result.py,sha256=Df_tPeqCQnLa0i0vVA-BGCJDx37ebD_3ojAmHnXE2yU,22767
 pydantic_ai/settings.py,sha256=ntuWnke9UA18aByDxk9OIhN0tAgOaPdqCEkRf-wlp8Y,3059
 pydantic_ai/tools.py,sha256=IPZuZJCSQUppz1uyLVwpfFLGoMirB8YtKWXIDQGR444,13414
-pydantic_ai/usage.py,sha256=60d9f6M7YEYuKMbqDGDogX4KsA73fhDtWyDXYXoIPaI,4948
+pydantic_ai/usage.py,sha256=VmpU_o_RjFI65J81G1wfCwDIAYBclMjeWfLtslntFOw,5406
 pydantic_ai/common_tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pydantic_ai/common_tools/duckduckgo.py,sha256=-kSa1gGn5-NIYvtxFWrFcX2XdmfEmGxI3_wAqrb6jLI,2230
 pydantic_ai/common_tools/tavily.py,sha256=Lz35037ggkdWKa_Stj0yXBkiN_hygDefEevoRDUclF0,2560
-pydantic_ai/models/__init__.py,sha256=Qw_g58KzGUmuDKOBa2u3yFrNbgCXkdRSNtkhseLC1VM,13758
+pydantic_ai/models/__init__.py,sha256=2A3CpdMnvllnVVX8PmlUcBs0HMGcG4RurOXsRKl0BPc,13886
 pydantic_ai/models/anthropic.py,sha256=bFtE6hku9L4l4pKJg8XER37T2ST2htArho5lPjEohAk,20637
 pydantic_ai/models/cohere.py,sha256=6F6eWPGVT7mpMXlRugbVbR-a8Q1zmb1SKS_fWOoBL80,11514
 pydantic_ai/models/fallback.py,sha256=smHwNIpxu19JsgYYjY0nmzl3yox7yQRJ0Ir08zdhnk0,4207
-pydantic_ai/models/function.py,sha256=EMlASu436RE-XzOTuHGkIqkS8J4WItUvwwaL08LLkX8,10948
+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=cvjHgQE_gJOH-YVQyvx9tBpGNB_Iuc8N8THn0TL0Rjk,8791
+pydantic_ai/models/instrumented.py,sha256=xUZEn2VG8hP3hny0L5kZgXC5UnFdlUJ0DgXOxFmYhEo,9654
 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.29.dist-info/METADATA,sha256=f7eLYpKWzEGmzTC8U29ZO5T4aTYMTW0Shmqa1zuEtG0,3062
-pydantic_ai_slim-0.0.29.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydantic_ai_slim-0.0.29.dist-info/RECORD,,
+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,,