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

pydantic_ai/durable_exec/temporal/_agent.py

@@ -17,6 +17,7 @@ from typing_extensions import Never
 
 from pydantic_ai import (
     AbstractToolset,
+    AgentRunResultEvent,
     _utils,
     messages as _messages,
     models,
@@ -258,7 +259,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -276,7 +277,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT],
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -293,7 +294,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT] | None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -368,7 +369,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -386,7 +387,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT],
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -403,7 +404,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT] | None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -476,7 +477,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -494,7 +495,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT],
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -512,7 +513,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT] | None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -558,9 +559,8 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         """
         if workflow.in_workflow():
             raise UserError(
-                '`agent.run_stream()` cannot currently be used inside a Temporal 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 Temporal workflow. '
+                'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
             )
 
         async with super().run_stream(
@@ -580,13 +580,131 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         ) as result:
             yield result
 
+    @overload
+    def run_stream_events(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: None = None,
+        message_history: Sequence[_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: Sequence[_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: Sequence[_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.
+        """
+        if workflow.in_workflow():
+            raise UserError(
+                '`agent.run_stream_events()` cannot be used inside a Temporal workflow. '
+                'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+            )
+
+        return super().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,
+        )
+
     @overload
     def iter(
         self,
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -604,7 +722,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT],
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -622,7 +740,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         user_prompt: str | Sequence[_messages.UserContent] | None = None,
         *,
         output_type: OutputSpec[RunOutputDataT] | None = None,
-        message_history: list[_messages.ModelMessage] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
         deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
@@ -711,9 +829,8 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         if workflow.in_workflow():
             if not self._temporal_overrides_active.get():
                 raise UserError(
-                    '`agent.iter()` cannot currently be used inside a Temporal 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.iter()` cannot be used inside a Temporal workflow. '
+                    'Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
                 )
 
             if model is not None:
@@ -745,18 +862,20 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
     def override(
         self,
         *,
+        name: str | _utils.Unset = _utils.UNSET,
         deps: AgentDepsT | _utils.Unset = _utils.UNSET,
         model: models.Model | models.KnownModelName | str | _utils.Unset = _utils.UNSET,
         toolsets: Sequence[AbstractToolset[AgentDepsT]] | _utils.Unset = _utils.UNSET,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] | _utils.Unset = _utils.UNSET,
         instructions: Instructions[AgentDepsT] | _utils.Unset = _utils.UNSET,
     ) -> Iterator[None]:
-        """Context manager to temporarily override agent dependencies, model, toolsets, tools, or instructions.
+        """Context manager to temporarily override agent name, dependencies, model, toolsets, tools, or instructions.
 
         This is particularly useful when testing.
         You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures).
 
         Args:
+            name: The name to use instead of the name passed to the agent constructor and agent run.
             deps: The dependencies to use instead of the dependencies passed to the agent run.
             model: The model to use instead of the model passed to the agent run.
             toolsets: The toolsets to use instead of the toolsets passed to the agent constructor and agent run.
@@ -778,6 +897,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
                 )
 
         with super().override(
+            name=name,
             deps=deps,
             model=model,
             toolsets=toolsets,

pydantic_ai/durable_exec/temporal/_model.py

@@ -128,6 +128,8 @@ class TemporalModel(WrapperModel):
         if not workflow.in_workflow():
             return await super().request(messages, model_settings, model_request_parameters)
 
+        self._validate_model_request_parameters(model_request_parameters)
+
         return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.request_activity,
             arg=_RequestParams(
@@ -163,6 +165,8 @@ class TemporalModel(WrapperModel):
         # and that only calls `request_stream` if `event_stream_handler` is set.
         assert self.event_stream_handler is not None
 
+        self._validate_model_request_parameters(model_request_parameters)
+
         serialized_run_context = self.run_context_type.serialize_run_context(run_context)
         response = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.request_stream_activity,
@@ -178,3 +182,7 @@ class TemporalModel(WrapperModel):
             **self.activity_config,
         )
         yield TemporalStreamedResponse(model_request_parameters, response)
+
+    def _validate_model_request_parameters(self, model_request_parameters: ModelRequestParameters) -> None:
+        if model_request_parameters.allow_image_output:
+            raise UserError('Image output is not supported with Temporal because of the 2MB payload size limit.')

pydantic_ai/format_prompt.py

@@ -2,7 +2,8 @@ from __future__ import annotations as _annotations
 
 from collections.abc import Iterable, Iterator, Mapping
 from dataclasses import asdict, dataclass, field, fields, is_dataclass
-from datetime import date
+from datetime import date, time, timedelta
+from enum import Enum
 from typing import Any, Literal
 from xml.etree import ElementTree
 
@@ -26,8 +27,8 @@ def format_as_xml(
     This is useful since LLMs often find it easier to read semi-structured data (e.g. examples) as XML,
     rather than JSON etc.
 
-    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `Mapping`,
-    `Iterable`, `dataclass`, and `BaseModel`.
+    Supports: `str`, `bytes`, `bytearray`, `bool`, `int`, `float`, `date`, `datetime`, `time`, `timedelta`, `Enum`,
+    `Mapping`, `Iterable`, `dataclass`, and `BaseModel`.
 
     Args:
         obj: Python Object to serialize to XML.
@@ -101,10 +102,12 @@ class _ToXml:
             element.text = value
         elif isinstance(value, bytes | bytearray):
             element.text = value.decode(errors='ignore')
-        elif isinstance(value, bool | int | float):
+        elif isinstance(value, bool | int | float | Enum):
             element.text = str(value)
-        elif isinstance(value, date):
+        elif isinstance(value, date | time):
             element.text = value.isoformat()
+        elif isinstance(value, timedelta):
+            element.text = str(value)
         elif isinstance(value, Mapping):
             if tag is None and path in self._element_names:
                 element.tag = self._element_names[path]
@@ -164,7 +167,7 @@ class _ToXml:
         path: str = '',
     ):
         """Parse data structures as dataclasses or Pydantic models to extract element names and attributes."""
-        if value is None or isinstance(value, (str | int | float | date | bytearray | bytes | bool)):
+        if value is None or isinstance(value, (str | int | float | date | time | timedelta | bytearray | bytes | bool)):
             return
         elif isinstance(value, Mapping):
             for k, v in value.items():  # pyright: ignore[reportUnknownVariableType]

pydantic_ai/mcp.py

@@ -167,6 +167,10 @@ class MCPServer(AbstractToolset[Any], ABC):
     def id(self) -> str | None:
         return self._id
 
+    @id.setter
+    def id(self, value: str | None):
+        self._id = value
+
     @property
     def label(self) -> str:
         if self.id:
@@ -414,6 +418,9 @@ class MCPServer(AbstractToolset[Any], ABC):
         else:
             assert_never(resource)
 
+    def __eq__(self, value: object, /) -> bool:
+        return isinstance(value, MCPServer) and self.id == value.id and self.tool_prefix == value.tool_prefix
+
 
 class MCPServerStdio(MCPServer):
     """Runs an MCP server in a subprocess and communicates with it over stdin/stdout.
@@ -568,10 +575,10 @@ class MCPServerStdio(MCPServer):
         return f'{self.__class__.__name__}({", ".join(repr_args)})'
 
     def __eq__(self, value: object, /) -> bool:
-        if not isinstance(value, MCPServerStdio):
-            return False  # pragma: no cover
         return (
-            self.command == value.command
+            super().__eq__(value)
+            and isinstance(value, MCPServerStdio)
+            and self.command == value.command
             and self.args == value.args
             and self.env == value.env
             and self.cwd == value.cwd
@@ -809,9 +816,7 @@ class MCPServerSSE(_MCPServerHTTP):
         return sse_client  # pragma: no cover
 
     def __eq__(self, value: object, /) -> bool:
-        if not isinstance(value, MCPServerSSE):
-            return False  # pragma: no cover
-        return self.url == value.url
+        return super().__eq__(value) and isinstance(value, MCPServerSSE) and self.url == value.url
 
 
 @deprecated('The `MCPServerHTTP` class is deprecated, use `MCPServerSSE` instead.')
@@ -885,9 +890,7 @@ class MCPServerStreamableHTTP(_MCPServerHTTP):
         return streamablehttp_client  # pragma: no cover
 
     def __eq__(self, value: object, /) -> bool:
-        if not isinstance(value, MCPServerStreamableHTTP):
-            return False  # pragma: no cover
-        return self.url == value.url
+        return super().__eq__(value) and isinstance(value, MCPServerStreamableHTTP) and self.url == value.url
 
 
 ToolResult = (
@@ -964,4 +967,11 @@ def load_mcp_servers(config_path: str | Path) -> list[MCPServerStdio | MCPServer
         raise FileNotFoundError(f'Config file {config_path} not found')
 
     config = MCPServerConfig.model_validate_json(config_path.read_bytes())
-    return list(config.mcp_servers.values())
+
+    servers: list[MCPServerStdio | MCPServerStreamableHTTP | MCPServerSSE] = []
+    for name, server in config.mcp_servers.items():
+        server.id = name
+        server.tool_prefix = name
+        servers.append(server)
+
+    return servers