pydantic-ai-slim 1.0.14__py3-none-any.whl → 1.0.15__py3-none-any.whl

pydantic_ai/_parts_manager.py

@@ -20,7 +20,6 @@ from typing import Any
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
     BuiltinToolCallPart,
-    BuiltinToolReturnPart,
     ModelResponsePart,
     ModelResponseStreamEvent,
     PartDeltaEvent,
@@ -350,64 +349,31 @@ class ModelResponsePartsManager:
             self._vendor_id_to_part_index[vendor_part_id] = new_part_index
         return PartStartEvent(index=new_part_index, part=new_part)
 
-    def handle_builtin_tool_call_part(
+    def handle_part(
         self,
         *,
         vendor_part_id: Hashable | None,
-        part: BuiltinToolCallPart,
+        part: ModelResponsePart,
     ) -> ModelResponseStreamEvent:
-        """Create or overwrite a BuiltinToolCallPart.
+        """Create or overwrite a ModelResponsePart.
 
         Args:
             vendor_part_id: The vendor's ID for this tool call part. If not
                 None and an existing part is found, that part is overwritten.
-            part: The BuiltinToolCallPart.
+            part: The ModelResponsePart.
 
         Returns:
-            ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part
-            has been added to the manager, or replaced an existing part.
-        """
-        if vendor_part_id is None:
-            # vendor_part_id is None, so we unconditionally append a new BuiltinToolCallPart to the end of the list
-            new_part_index = len(self._parts)
-            self._parts.append(part)
-        else:
-            # vendor_part_id is provided, so find and overwrite or create a new BuiltinToolCallPart.
-            maybe_part_index = self._vendor_id_to_part_index.get(vendor_part_id)
-            if maybe_part_index is not None and isinstance(self._parts[maybe_part_index], BuiltinToolCallPart):
-                new_part_index = maybe_part_index
-                self._parts[new_part_index] = part
-            else:
-                new_part_index = len(self._parts)
-                self._parts.append(part)
-            self._vendor_id_to_part_index[vendor_part_id] = new_part_index
-        return PartStartEvent(index=new_part_index, part=part)
-
-    def handle_builtin_tool_return_part(
-        self,
-        *,
-        vendor_part_id: Hashable | None,
-        part: BuiltinToolReturnPart,
-    ) -> ModelResponseStreamEvent:
-        """Create or overwrite a BuiltinToolReturnPart.
-
-        Args:
-            vendor_part_id: The vendor's ID for this tool call part. If not
-                None and an existing part is found, that part is overwritten.
-            part: The BuiltinToolReturnPart.
-
-        Returns:
-            ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part
+            ModelResponseStreamEvent: A `PartStartEvent` indicating that a new part
             has been added to the manager, or replaced an existing part.
         """
         if vendor_part_id is None:
-            # vendor_part_id is None, so we unconditionally append a new BuiltinToolReturnPart to the end of the list
+            # vendor_part_id is None, so we unconditionally append a new part to the end of the list
             new_part_index = len(self._parts)
             self._parts.append(part)
         else:
-            # vendor_part_id is provided, so find and overwrite or create a new BuiltinToolReturnPart.
+            # vendor_part_id is provided, so find and overwrite or create a new part.
             maybe_part_index = self._vendor_id_to_part_index.get(vendor_part_id)
-            if maybe_part_index is not None and isinstance(self._parts[maybe_part_index], BuiltinToolReturnPart):
+            if maybe_part_index is not None and isinstance(self._parts[maybe_part_index], type(part)):
                 new_part_index = maybe_part_index
                 self._parts[new_part_index] = part
             else:

pydantic_ai/_tool_manager.py

@@ -18,7 +18,7 @@ from .exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior
 from .messages import ToolCallPart
 from .tools import ToolDefinition
 from .toolsets.abstract import AbstractToolset, ToolsetTool
-from .usage import UsageLimits
+from .usage import RunUsage
 
 _sequential_tool_calls_ctx_var: ContextVar[bool] = ContextVar('sequential_tool_calls', default=False)
 
@@ -93,7 +93,6 @@ class ToolManager(Generic[AgentDepsT]):
         call: ToolCallPart,
         allow_partial: bool = False,
         wrap_validation_errors: bool = True,
-        usage_limits: UsageLimits | None = None,
     ) -> Any:
         """Handle a tool call by validating the arguments, calling the tool, and handling retries.
 
@@ -108,16 +107,16 @@ class ToolManager(Generic[AgentDepsT]):
 
         if (tool := self.tools.get(call.tool_name)) and tool.tool_def.kind == 'output':
             # Output tool calls are not traced and not counted
-            return await self._call_tool(call, allow_partial, wrap_validation_errors, count_tool_usage=False)
+            return await self._call_tool(call, allow_partial, wrap_validation_errors)
         else:
-            return await self._call_tool_traced(
+            return await self._call_function_tool(
                 call,
                 allow_partial,
                 wrap_validation_errors,
                 self.ctx.tracer,
                 self.ctx.trace_include_content,
                 self.ctx.instrumentation_version,
-                usage_limits,
+                self.ctx.usage,
             )
 
     async def _call_tool(
@@ -125,8 +124,6 @@ class ToolManager(Generic[AgentDepsT]):
         call: ToolCallPart,
         allow_partial: bool,
         wrap_validation_errors: bool,
-        usage_limits: UsageLimits | None = None,
-        count_tool_usage: bool = True,
     ) -> Any:
         if self.tools is None or self.ctx is None:
             raise ValueError('ToolManager has not been prepared for a run step yet')  # pragma: no cover
@@ -159,14 +156,8 @@ class ToolManager(Generic[AgentDepsT]):
             else:
                 args_dict = validator.validate_python(call.args or {}, allow_partial=pyd_allow_partial)
 
-            if usage_limits is not None and count_tool_usage:
-                usage_limits.check_before_tool_call(self.ctx.usage)
-
             result = await self.toolset.call_tool(name, args_dict, ctx, tool)
 
-            if count_tool_usage:
-                self.ctx.usage.tool_calls += 1
-
             return result
         except (ValidationError, ModelRetry) as e:
             max_retries = tool.max_retries if tool is not None else 1
@@ -199,7 +190,7 @@ class ToolManager(Generic[AgentDepsT]):
 
                 raise e
 
-    async def _call_tool_traced(
+    async def _call_function_tool(
         self,
         call: ToolCallPart,
         allow_partial: bool,
@@ -207,7 +198,7 @@ class ToolManager(Generic[AgentDepsT]):
         tracer: Tracer,
         include_content: bool,
         instrumentation_version: int,
-        usage_limits: UsageLimits | None = None,
+        usage: RunUsage,
     ) -> Any:
         """See <https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#execute-tool-span>."""
         instrumentation_names = InstrumentationNames.for_version(instrumentation_version)
@@ -242,7 +233,9 @@ class ToolManager(Generic[AgentDepsT]):
             attributes=span_attributes,
         ) as span:
             try:
-                tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors, usage_limits)
+                tool_result = await self._call_tool(call, allow_partial, wrap_validation_errors)
+                usage.tool_calls += 1
+
             except ToolRetryError as e:
                 part = e.tool_retry
                 if include_content and span.is_recording():

pydantic_ai/agent/abstract.py

@@ -1,5 +1,6 @@
 from __future__ import annotations as _annotations
 
+import asyncio
 import inspect
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
@@ -7,6 +8,7 @@ from contextlib import AbstractAsyncContextManager, asynccontextmanager, context
 from types import FrameType
 from typing import TYPE_CHECKING, Any, Generic, TypeAlias, cast, overload
 
+import anyio
 from typing_extensions import Self, TypeIs, TypeVar
 
 from pydantic_graph import End
@@ -25,7 +27,7 @@ from .. import (
 from .._tool_manager import ToolManager
 from ..output import OutputDataT, OutputSpec
 from ..result import AgentStream, FinalResult, StreamedRunResult
-from ..run import AgentRun, AgentRunResult
+from ..run import AgentRun, AgentRunResult, AgentRunResultEvent
 from ..settings import ModelSettings
 from ..tools import (
     AgentDepsT,
@@ -552,6 +554,172 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         if not yielded:
             raise exceptions.AgentRunError('Agent run finished without producing a final result')  # pragma: no cover
 
+    @overload
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[OutputDataT]]: ...
+
+    @overload
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT],
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[RunOutputDataT]]: ...
+
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[Any]]:
+        """Run the agent with a user prompt in async mode and stream events from the run.
+
+        This is a convenience method that wraps [`self.run`][pydantic_ai.agent.AbstractAgent.run] and
+        uses the `event_stream_handler` kwarg to get a stream of events from the run.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent, AgentRunResultEvent, AgentStreamEvent
+
+        agent = Agent('openai:gpt-4o')
+
+        async def main():
+            events: list[AgentStreamEvent | AgentRunResultEvent] = []
+            async for event in agent.run_stream_events('What is the capital of France?'):
+                events.append(event)
+            print(events)
+            '''
+            [
+                PartStartEvent(index=0, part=TextPart(content='The capital of ')),
+                FinalResultEvent(tool_name=None, tool_call_id=None),
+                PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='France is Paris. ')),
+                AgentRunResultEvent(
+                    result=AgentRunResult(output='The capital of France is Paris. ')
+                ),
+            ]
+            '''
+        ```
+
+        Arguments are the same as for [`self.run`][pydantic_ai.agent.AbstractAgent.run],
+        except that `event_stream_handler` is now allowed.
+
+        Args:
+            user_prompt: User input to start/continue the conversation.
+            output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
+                output validators since output validators would expect an argument that matches the agent's output type.
+            message_history: History of the conversation so far.
+            deferred_tool_results: Optional results for deferred tool calls in the message history.
+            model: Optional model to use for this run, required if `model` was not set when creating the agent.
+            deps: Optional dependencies to use for this run.
+            model_settings: Optional settings to use for this model's request.
+            usage_limits: Optional limits on model request count or token usage.
+            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
+            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
+            toolsets: Optional additional toolsets for this run.
+
+        Returns:
+            An async iterable of stream events `AgentStreamEvent` and finally a `AgentRunResultEvent` with the final
+            run result.
+        """
+        # unfortunately this hack of returning a generator rather than defining it right here is
+        # required to allow overloads of this method to work in python's typing system, or at least with pyright
+        # or at least I couldn't make it work without
+        return self._run_stream_events(
+            user_prompt,
+            output_type=output_type,
+            message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
+            model=model,
+            deps=deps,
+            model_settings=model_settings,
+            usage_limits=usage_limits,
+            usage=usage,
+            infer_name=infer_name,
+            toolsets=toolsets,
+        )
+
+    async def _run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[Any]]:
+        send_stream, receive_stream = anyio.create_memory_object_stream[
+            _messages.AgentStreamEvent | AgentRunResultEvent[Any]
+        ]()
+
+        async def event_stream_handler(
+            _: RunContext[AgentDepsT], events: AsyncIterable[_messages.AgentStreamEvent]
+        ) -> None:
+            async for event in events:
+                await send_stream.send(event)
+
+        async def run_agent() -> AgentRunResult[Any]:
+            async with send_stream:
+                return await self.run(
+                    user_prompt,
+                    output_type=output_type,
+                    message_history=message_history,
+                    deferred_tool_results=deferred_tool_results,
+                    model=model,
+                    deps=deps,
+                    model_settings=model_settings,
+                    usage_limits=usage_limits,
+                    usage=usage,
+                    infer_name=infer_name,
+                    toolsets=toolsets,
+                    event_stream_handler=event_stream_handler,
+                )
+
+        task = asyncio.create_task(run_agent())
+
+        async with receive_stream:
+            async for message in receive_stream:
+                yield message
+
+        result = await task
+        yield AgentRunResultEvent(result)
+
     @overload
     def iter(
         self,

pydantic_ai/builtin_tools.py

@@ -12,6 +12,7 @@ __all__ = (
     'WebSearchUserLocation',
     'CodeExecutionTool',
     'UrlContextTool',
+    'ImageGenerationTool',
     'MemoryTool',
 )
 
@@ -142,6 +143,87 @@ class UrlContextTool(AbstractBuiltinTool):
     """The kind of tool."""
 
 
+@dataclass(kw_only=True)
+class ImageGenerationTool(AbstractBuiltinTool):
+    """A builtin tool that allows your agent to generate images.
+
+    Supported by:
+
+    * OpenAI Responses
+    * Google
+    """
+
+    background: Literal['transparent', 'opaque', 'auto'] = 'auto'
+    """Background type for the generated image.
+
+    Supported by:
+
+    * OpenAI Responses. 'transparent' is only supported for 'png' and 'webp' output formats.
+    """
+
+    input_fidelity: Literal['high', 'low'] | None = None
+    """
+    Control how much effort the model will exert to match the style and features,
+    especially facial features, of input images.
+
+    Supported by:
+
+    * OpenAI Responses. Default: 'low'.
+    """
+
+    moderation: Literal['auto', 'low'] = 'auto'
+    """Moderation level for the generated image.
+
+    Supported by:
+
+    * OpenAI Responses
+    """
+
+    output_compression: int = 100
+    """Compression level for the output image.
+
+    Supported by:
+
+    * OpenAI Responses. Only supported for 'png' and 'webp' output formats.
+    """
+
+    output_format: Literal['png', 'webp', 'jpeg'] | None = None
+    """The output format of the generated image.
+
+    Supported by:
+
+    * OpenAI Responses. Default: 'png'.
+    """
+
+    partial_images: int = 0
+    """
+    Number of partial images to generate in streaming mode.
+
+    Supported by:
+
+    * OpenAI Responses. Supports 0 to 3.
+    """
+
+    quality: Literal['low', 'medium', 'high', 'auto'] = 'auto'
+    """The quality of the generated image.
+
+    Supported by:
+
+    * OpenAI Responses
+    """
+
+    size: Literal['1024x1024', '1024x1536', '1536x1024', 'auto'] = 'auto'
+    """The size of the generated image.
+
+    Supported by:
+
+    * OpenAI Responses
+    """
+
+    kind: str = 'image_generation'
+    """The kind of tool."""
+
+
 class MemoryTool(AbstractBuiltinTool):
     """A builtin tool that allows your agent to use memory.
 

pydantic_ai/direct.py

@@ -364,10 +364,17 @@ class StreamedResponseSync:
         if self._thread and self._thread.is_alive():
             self._thread.join()
 
+    # TODO (v2): Drop in favor of `response` property
     def get(self) -> messages.ModelResponse:
         """Build a ModelResponse from the data received from the stream so far."""
         return self._ensure_stream_ready().get()
 
+    @property
+    def response(self) -> messages.ModelResponse:
+        """Get the current state of the response."""
+        return self.get()
+
+    # TODO (v2): Make this a property
     def usage(self) -> RequestUsage:
         """Get the usage of the response so far."""
         return self._ensure_stream_ready().usage()

pydantic_ai/durable_exec/dbos/_agent.py

@@ -9,6 +9,7 @@ from typing_extensions import Never
 
 from pydantic_ai import (
     AbstractToolset,
+    AgentRunResultEvent,
     _utils,
     messages as _messages,
     models,
@@ -525,9 +526,8 @@ class DBOSAgent(WrapperAgent[AgentDepsT, OutputDataT], DBOSConfiguredInstance):
         """
         if DBOS.workflow_id is not None and DBOS.step_id is None:
             raise UserError(
-                '`agent.run_stream()` cannot currently be used inside a DBOS workflow. '
-                'Set an `event_stream_handler` on the agent and use `agent.run()` instead. '
-                'Please file an issue if this is not sufficient for your use case.'
+                '`agent.run_stream()` cannot be used inside a DBOS workflow. '
+                'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
             )
 
         async with super().run_stream(
@@ -547,6 +547,109 @@ class DBOSAgent(WrapperAgent[AgentDepsT, OutputDataT], DBOSConfiguredInstance):
         ) as result:
             yield result
 
+    @overload
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[OutputDataT]]: ...
+
+    @overload
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT],
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[RunOutputDataT]]: ...
+
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
+        model: models.Model | models.KnownModelName | str | None = None,
+        deps: AgentDepsT = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.RunUsage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[_messages.AgentStreamEvent | AgentRunResultEvent[Any]]:
+        """Run the agent with a user prompt in async mode and stream events from the run.
+
+        This is a convenience method that wraps [`self.run`][pydantic_ai.agent.AbstractAgent.run] and
+        uses the `event_stream_handler` kwarg to get a stream of events from the run.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent, AgentRunResultEvent, AgentStreamEvent
+
+        agent = Agent('openai:gpt-4o')
+
+        async def main():
+            events: list[AgentStreamEvent | AgentRunResultEvent] = []
+            async for event in agent.run_stream_events('What is the capital of France?'):
+                events.append(event)
+            print(events)
+            '''
+            [
+                PartStartEvent(index=0, part=TextPart(content='The capital of ')),
+                FinalResultEvent(tool_name=None, tool_call_id=None),
+                PartDeltaEvent(index=0, delta=TextPartDelta(content_delta='France is Paris. ')),
+                AgentRunResultEvent(
+                    result=AgentRunResult(output='The capital of France is Paris. ')
+                ),
+            ]
+            '''
+        ```
+
+        Arguments are the same as for [`self.run`][pydantic_ai.agent.AbstractAgent.run],
+        except that `event_stream_handler` is now allowed.
+
+        Args:
+            user_prompt: User input to start/continue the conversation.
+            output_type: Custom output type to use for this run, `output_type` may only be used if the agent has no
+                output validators since output validators would expect an argument that matches the agent's output type.
+            message_history: History of the conversation so far.
+            deferred_tool_results: Optional results for deferred tool calls in the message history.
+            model: Optional model to use for this run, required if `model` was not set when creating the agent.
+            deps: Optional dependencies to use for this run.
+            model_settings: Optional settings to use for this model's request.
+            usage_limits: Optional limits on model request count or token usage.
+            usage: Optional usage to start with, useful for resuming a conversation or agents used in tools.
+            infer_name: Whether to try to infer the agent name from the call frame if it's not set.
+            toolsets: Optional additional toolsets for this run.
+
+        Returns:
+            An async iterable of stream events `AgentStreamEvent` and finally a `AgentRunResultEvent` with the final
+            run result.
+        """
+        raise UserError(
+            '`agent.run_stream_events()` cannot be used with DBOS. '
+            'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+        )
+
     @overload
     def iter(
         self,