pydantic-ai-slim 1.9.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

pydantic_ai/durable_exec/temporal/_function_toolset.py

@@ -1,57 +1,22 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Annotated, Any, Literal
+from typing import Any, Literal
 
-from pydantic import ConfigDict, Discriminator, with_config
 from temporalio import activity, workflow
 from temporalio.workflow import ActivityConfig
-from typing_extensions import assert_never
 
 from pydantic_ai import FunctionToolset, ToolsetTool
-from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.tools import AgentDepsT, RunContext
 from pydantic_ai.toolsets.function import FunctionToolsetTool
 
 from ._run_context import TemporalRunContext
-from ._toolset import TemporalWrapperToolset
-
-
-@dataclass
-@with_config(ConfigDict(arbitrary_types_allowed=True))
-class _CallToolParams:
-    name: str
-    tool_args: dict[str, Any]
-    serialized_run_context: Any
-
-
-@dataclass
-class _ApprovalRequired:
-    kind: Literal['approval_required'] = 'approval_required'
-
-
-@dataclass
-class _CallDeferred:
-    kind: Literal['call_deferred'] = 'call_deferred'
-
-
-@dataclass
-class _ModelRetry:
-    message: str
-    kind: Literal['model_retry'] = 'model_retry'
-
-
-@dataclass
-class _ToolReturn:
-    result: Any
-    kind: Literal['tool_return'] = 'tool_return'
-
-
-_CallToolResult = Annotated[
-    _ApprovalRequired | _CallDeferred | _ModelRetry | _ToolReturn,
-    Discriminator('kind'),
-]
+from ._toolset import (
+    CallToolParams,
+    CallToolResult,
+    TemporalWrapperToolset,
+)
 
 
 class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
@@ -70,7 +35,7 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
         self.tool_activity_config = tool_activity_config
         self.run_context_type = run_context_type
 
-        async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _CallToolResult:
+        async def call_tool_activity(params: CallToolParams, deps: AgentDepsT) -> CallToolResult:
             name = params.name
             ctx = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             try:
@@ -84,15 +49,7 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
             # The tool args will already have been validated into their proper types in the `ToolManager`,
             # but `execute_activity` would have turned them into simple Python types again, so we need to re-validate them.
             args_dict = tool.args_validator.validate_python(params.tool_args)
-            try:
-                result = await self.wrapped.call_tool(name, args_dict, ctx, tool)
-                return _ToolReturn(result=result)
-            except ApprovalRequired:
-                return _ApprovalRequired()
-            except CallDeferred:
-                return _CallDeferred()
-            except ModelRetry as e:
-                return _ModelRetry(message=e.message)
+            return await self._wrap_call_tool_result(self.wrapped.call_tool(name, args_dict, ctx, tool))
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         call_tool_activity.__annotations__['deps'] = deps_type
@@ -123,25 +80,18 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
 
         tool_activity_config = self.activity_config | tool_activity_config
         serialized_run_context = self.run_context_type.serialize_run_context(ctx)
-        result = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
-            activity=self.call_tool_activity,
-            args=[
-                _CallToolParams(
-                    name=name,
-                    tool_args=tool_args,
-                    serialized_run_context=serialized_run_context,
-                ),
-                ctx.deps,
-            ],
-            **tool_activity_config,
+        return self._unwrap_call_tool_result(
+            await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+                activity=self.call_tool_activity,
+                args=[
+                    CallToolParams(
+                        name=name,
+                        tool_args=tool_args,
+                        serialized_run_context=serialized_run_context,
+                        tool_def=None,
+                    ),
+                    ctx.deps,
+                ],
+                **tool_activity_config,
+            )
         )
-        if isinstance(result, _ApprovalRequired):
-            raise ApprovalRequired()
-        elif isinstance(result, _CallDeferred):
-            raise CallDeferred()
-        elif isinstance(result, _ModelRetry):
-            raise ModelRetry(result.message)
-        elif isinstance(result, _ToolReturn):
-            return result.result
-        else:
-            assert_never(result)

pydantic_ai/durable_exec/temporal/_mcp_server.py

@@ -11,11 +11,15 @@ from typing_extensions import Self
 
 from pydantic_ai import ToolsetTool
 from pydantic_ai.exceptions import UserError
-from pydantic_ai.mcp import MCPServer, ToolResult
+from pydantic_ai.mcp import MCPServer
 from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
 
 from ._run_context import TemporalRunContext
-from ._toolset import TemporalWrapperToolset
+from ._toolset import (
+    CallToolParams,
+    CallToolResult,
+    TemporalWrapperToolset,
+)
 
 
 @dataclass
@@ -24,15 +28,6 @@ class _GetToolsParams:
     serialized_run_context: Any
 
 
-@dataclass
-@with_config(ConfigDict(arbitrary_types_allowed=True))
-class _CallToolParams:
-    name: str
-    tool_args: dict[str, Any]
-    serialized_run_context: Any
-    tool_def: ToolDefinition
-
-
 class TemporalMCPServer(TemporalWrapperToolset[AgentDepsT]):
     def __init__(
         self,
@@ -72,13 +67,16 @@ class TemporalMCPServer(TemporalWrapperToolset[AgentDepsT]):
             get_tools_activity
         )
 
-        async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> ToolResult:
+        async def call_tool_activity(params: CallToolParams, deps: AgentDepsT) -> CallToolResult:
             run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
-            return await self.wrapped.call_tool(
-                params.name,
-                params.tool_args,
-                run_context,
-                self.tool_for_tool_def(params.tool_def),
+            assert isinstance(params.tool_def, ToolDefinition)
+            return await self._wrap_call_tool_result(
+                self.wrapped.call_tool(
+                    params.name,
+                    params.tool_args,
+                    run_context,
+                    self.tool_for_tool_def(params.tool_def),
+                )
             )
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
@@ -125,22 +123,24 @@ class TemporalMCPServer(TemporalWrapperToolset[AgentDepsT]):
         tool_args: dict[str, Any],
         ctx: RunContext[AgentDepsT],
         tool: ToolsetTool[AgentDepsT],
-    ) -> ToolResult:
+    ) -> CallToolResult:
         if not workflow.in_workflow():
             return await super().call_tool(name, tool_args, ctx, tool)
 
         tool_activity_config = self.activity_config | self.tool_activity_config.get(name, {})
         serialized_run_context = self.run_context_type.serialize_run_context(ctx)
-        return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
-            activity=self.call_tool_activity,
-            args=[
-                _CallToolParams(
-                    name=name,
-                    tool_args=tool_args,
-                    serialized_run_context=serialized_run_context,
-                    tool_def=tool.tool_def,
-                ),
-                ctx.deps,
-            ],
-            **tool_activity_config,
+        return self._unwrap_call_tool_result(
+            await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+                activity=self.call_tool_activity,
+                args=[
+                    CallToolParams(
+                        name=name,
+                        tool_args=tool_args,
+                        serialized_run_context=serialized_run_context,
+                        tool_def=tool.tool_def,
+                    ),
+                    ctx.deps,
+                ],
+                **tool_activity_config,
+            )
         )

pydantic_ai/durable_exec/temporal/_run_context.py

@@ -2,14 +2,19 @@ from __future__ import annotations
 
 from typing import Any
 
+from typing_extensions import TypeVar
+
 from pydantic_ai.exceptions import UserError
-from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.tools import RunContext
+
+AgentDepsT = TypeVar('AgentDepsT', default=None, covariant=True)
+"""Type variable for the agent dependencies in `RunContext`."""
 
 
 class TemporalRunContext(RunContext[AgentDepsT]):
     """The [`RunContext`][pydantic_ai.tools.RunContext] subclass to use to serialize and deserialize the run context for use inside a Temporal activity.
 
-    By default, only the `deps`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries` and `run_step` attributes will be available.
+    By default, only the `deps`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries`, `run_step` and `partial_output` attributes will be available.
     To make another attribute available, create a `TemporalRunContext` subclass with a custom `serialize_run_context` class method that returns a dictionary that includes the attribute and pass it to [`TemporalAgent`][pydantic_ai.durable_exec.temporal.TemporalAgent].
     """
 
@@ -44,9 +49,10 @@ class TemporalRunContext(RunContext[AgentDepsT]):
             'retry': ctx.retry,
             'max_retries': ctx.max_retries,
             'run_step': ctx.run_step,
+            'partial_output': ctx.partial_output,
         }
 
     @classmethod
-    def deserialize_run_context(cls, ctx: dict[str, Any], deps: AgentDepsT) -> TemporalRunContext[AgentDepsT]:
+    def deserialize_run_context(cls, ctx: dict[str, Any], deps: Any) -> TemporalRunContext[Any]:
         """Deserialize the run context from a `dict[str, Any]`."""
         return cls(**ctx, deps=deps)

pydantic_ai/durable_exec/temporal/_toolset.py

@@ -1,17 +1,58 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
-from typing import Any, Literal
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from typing import Annotated, Any, Literal
 
+from pydantic import ConfigDict, Discriminator, with_config
 from temporalio.workflow import ActivityConfig
+from typing_extensions import assert_never
 
 from pydantic_ai import AbstractToolset, FunctionToolset, WrapperToolset
-from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry
+from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
 from ._run_context import TemporalRunContext
 
 
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class CallToolParams:
+    name: str
+    tool_args: dict[str, Any]
+    serialized_run_context: Any
+    tool_def: ToolDefinition | None
+
+
+@dataclass
+class _ApprovalRequired:
+    kind: Literal['approval_required'] = 'approval_required'
+
+
+@dataclass
+class _CallDeferred:
+    kind: Literal['call_deferred'] = 'call_deferred'
+
+
+@dataclass
+class _ModelRetry:
+    message: str
+    kind: Literal['model_retry'] = 'model_retry'
+
+
+@dataclass
+class _ToolReturn:
+    result: Any
+    kind: Literal['tool_return'] = 'tool_return'
+
+
+CallToolResult = Annotated[
+    _ApprovalRequired | _CallDeferred | _ModelRetry | _ToolReturn,
+    Discriminator('kind'),
+]
+
+
 class TemporalWrapperToolset(WrapperToolset[AgentDepsT], ABC):
     @property
     def id(self) -> str:
@@ -30,6 +71,29 @@ class TemporalWrapperToolset(WrapperToolset[AgentDepsT], ABC):
         # Temporalized toolsets cannot be swapped out after the fact.
         return self
 
+    async def _wrap_call_tool_result(self, coro: Awaitable[Any]) -> CallToolResult:
+        try:
+            result = await coro
+            return _ToolReturn(result=result)
+        except ApprovalRequired:
+            return _ApprovalRequired()
+        except CallDeferred:
+            return _CallDeferred()
+        except ModelRetry as e:
+            return _ModelRetry(message=e.message)
+
+    def _unwrap_call_tool_result(self, result: CallToolResult) -> Any:
+        if isinstance(result, _ToolReturn):
+            return result.result
+        elif isinstance(result, _ApprovalRequired):
+            raise ApprovalRequired()
+        elif isinstance(result, _CallDeferred):
+            raise CallDeferred()
+        elif isinstance(result, _ModelRetry):
+            raise ModelRetry(result.message)
+        else:
+            assert_never(result)
+
 
 def temporalize_toolset(
     toolset: AbstractToolset[AgentDepsT],

pydantic_ai/mcp.py

@@ -726,9 +726,9 @@ class _MCPServerHTTP(MCPServer):
             MemoryObjectReceiveStream[SessionMessage | Exception],
             MemoryObjectSendStream[SessionMessage],
         ]
-    ]:  # pragma: no cover
+    ]:
         if self.http_client and self.headers:
-            raise ValueError('`http_client` is mutually exclusive with `headers`.')
+            raise ValueError('`http_client` is mutually exclusive with `headers`.')  # pragma: no cover
 
         transport_client_partial = functools.partial(
             self._transport_client,
@@ -737,7 +737,7 @@ class _MCPServerHTTP(MCPServer):
             sse_read_timeout=self.read_timeout,
         )
 
-        if self.http_client is not None:
+        if self.http_client is not None:  # pragma: no cover
 
             def httpx_client_factory(
                 headers: dict[str, str] | None = None,
@@ -866,7 +866,7 @@ class MCPServerStreamableHTTP(_MCPServerHTTP):
 
     @property
     def _transport_client(self):
-        return streamablehttp_client  # pragma: no cover
+        return streamablehttp_client
 
     def __eq__(self, value: object, /) -> bool:
         return super().__eq__(value) and isinstance(value, MCPServerStreamableHTTP) and self.url == value.url

pydantic_ai/messages.py

@@ -34,6 +34,7 @@ DocumentMediaType: TypeAlias = Literal[
     'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
     'text/html',
     'text/markdown',
+    'application/msword',
     'application/vnd.ms-excel',
 ]
 VideoMediaType: TypeAlias = Literal[
@@ -434,8 +435,12 @@ class DocumentUrl(FileUrl):
             return 'application/pdf'
         elif self.url.endswith('.rtf'):
             return 'application/rtf'
+        elif self.url.endswith('.doc'):
+            return 'application/msword'
         elif self.url.endswith('.docx'):
             return 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
+        elif self.url.endswith('.xls'):
+            return 'application/vnd.ms-excel'
         elif self.url.endswith('.xlsx'):
             return 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
 
@@ -480,7 +485,7 @@ class BinaryContent:
     """
 
     _identifier: Annotated[str | None, pydantic.Field(alias='identifier', default=None, exclude=True)] = field(
-        compare=False, default=None, repr=False
+        compare=False, default=None
     )
 
     kind: Literal['binary'] = 'binary'
@@ -645,6 +650,7 @@ _document_format_lookup: dict[str, DocumentFormat] = {
     'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': 'xlsx',
     'text/html': 'html',
     'text/markdown': 'md',
+    'application/msword': 'doc',
     'application/vnd.ms-excel': 'xls',
 }
 _audio_format_lookup: dict[str, AudioFormat] = {
@@ -882,7 +888,10 @@ class RetryPromptPart:
                 description = self.content
         else:
             json_errors = error_details_ta.dump_json(self.content, exclude={'__all__': {'ctx'}}, indent=2)
-            description = f'{len(self.content)} validation errors: {json_errors.decode()}'
+            plural = isinstance(self.content, list) and len(self.content) != 1
+            description = (
+                f'{len(self.content)} validation error{"s" if plural else ""}:\n```json\n{json_errors.decode()}\n```'
+            )
         return f'{description}\n\nFix the errors and try again.'
 
     def otel_event(self, settings: InstrumentationSettings) -> Event:

pydantic_ai/models/__init__.py

@@ -21,7 +21,7 @@ from typing_extensions import TypeAliasType, TypedDict
 
 from .. import _utils
 from .._json_schema import JsonSchemaTransformer
-from .._output import OutputObjectDefinition
+from .._output import OutputObjectDefinition, PromptedOutputSchema
 from .._parts_manager import ModelResponsePartsManager
 from .._run_context import RunContext
 from ..builtin_tools import AbstractBuiltinTool
@@ -309,6 +309,7 @@ class ModelRequestParameters:
     output_mode: OutputMode = 'text'
     output_object: OutputObjectDefinition | None = None
     output_tools: list[ToolDefinition] = field(default_factory=list)
+    prompted_output_template: str | None = None
     allow_text_output: bool = True
     allow_image_output: bool = False
 
@@ -316,6 +317,12 @@ class ModelRequestParameters:
     def tool_defs(self) -> dict[str, ToolDefinition]:
         return {tool_def.name: tool_def for tool_def in [*self.function_tools, *self.output_tools]}
 
+    @cached_property
+    def prompted_output_instructions(self) -> str | None:
+        if self.output_mode == 'prompted' and self.prompted_output_template and self.output_object:
+            return PromptedOutputSchema.build_instructions(self.prompted_output_template, self.output_object)
+        return None
+
     __repr__ = _utils.dataclasses_no_defaults_repr
 
 
@@ -408,23 +415,52 @@ class Model(ABC):
     ) -> tuple[ModelSettings | None, ModelRequestParameters]:
         """Prepare request inputs before they are passed to the provider.
 
-        This merges the given ``model_settings`` with the model's own ``settings`` attribute and ensures
-        ``customize_request_parameters`` is applied to the resolved
+        This merges the given `model_settings` with the model's own `settings` attribute and ensures
+        `customize_request_parameters` is applied to the resolved
         [`ModelRequestParameters`][pydantic_ai.models.ModelRequestParameters]. Subclasses can override this method if
         they need to customize the preparation flow further, but most implementations should simply call
-        ``self.prepare_request(...)`` at the start of their ``request`` (and related) methods.
+        `self.prepare_request(...)` at the start of their `request` (and related) methods.
         """
         model_settings = merge_model_settings(self.settings, model_settings)
 
-        if builtin_tools := model_request_parameters.builtin_tools:
+        params = self.customize_request_parameters(model_request_parameters)
+
+        if builtin_tools := params.builtin_tools:
             # Deduplicate builtin tools
-            model_request_parameters = replace(
-                model_request_parameters,
+            params = replace(
+                params,
                 builtin_tools=list({tool.unique_id: tool for tool in builtin_tools}.values()),
             )
 
-        model_request_parameters = self.customize_request_parameters(model_request_parameters)
-        return model_settings, model_request_parameters
+        if params.output_mode == 'auto':
+            output_mode = self.profile.default_structured_output_mode
+            params = replace(
+                params,
+                output_mode=output_mode,
+                allow_text_output=output_mode in ('native', 'prompted'),
+            )
+
+        # Reset irrelevant fields
+        if params.output_tools and params.output_mode != 'tool':
+            params = replace(params, output_tools=[])
+        if params.output_object and params.output_mode not in ('native', 'prompted'):
+            params = replace(params, output_object=None)
+        if params.prompted_output_template and params.output_mode != 'prompted':
+            params = replace(params, prompted_output_template=None)  # pragma: no cover
+
+        # Set default prompted output template
+        if params.output_mode == 'prompted' and not params.prompted_output_template:
+            params = replace(params, prompted_output_template=self.profile.prompted_output_template)
+
+        # Check if output mode is supported
+        if params.output_mode == 'native' and not self.profile.supports_json_schema_output:
+            raise UserError('Native structured output is not supported by this model.')
+        if params.output_mode == 'tool' and not self.profile.supports_tools:
+            raise UserError('Tool output is not supported by this model.')
+        if params.allow_image_output and not self.profile.supports_image_output:
+            raise UserError('Image output is not supported by this model.')
+
+        return model_settings, params
 
     @property
     @abstractmethod
@@ -462,13 +498,17 @@ class Model(ABC):
         return None
 
     @staticmethod
-    def _get_instructions(messages: list[ModelMessage]) -> str | None:
+    def _get_instructions(
+        messages: list[ModelMessage], model_request_parameters: ModelRequestParameters | None = None
+    ) -> str | None:
         """Get instructions from the first ModelRequest found when iterating messages in reverse.
 
         In the case that a "mock" request was generated to include a tool-return part for a result tool,
         we want to use the instructions from the second-to-most-recent request (which should correspond to the
         original request that generated the response that resulted in the tool-return part).
         """
+        instructions = None
+
         last_two_requests: list[ModelRequest] = []
         for message in reversed(messages):
             if isinstance(message, ModelRequest):
@@ -476,33 +516,38 @@ class Model(ABC):
                 if len(last_two_requests) == 2:
                     break
                 if message.instructions is not None:
-                    return message.instructions
+                    instructions = message.instructions
+                    break
 
         # If we don't have two requests, and we didn't already return instructions, there are definitely not any:
-        if len(last_two_requests) != 2:
-            return None
-
-        most_recent_request = last_two_requests[0]
-        second_most_recent_request = last_two_requests[1]
-
-        # If we've gotten this far and the most recent request consists of only tool-return parts or retry-prompt parts,
-        # we use the instructions from the second-to-most-recent request. This is necessary because when handling
-        # result tools, we generate a "mock" ModelRequest with a tool-return part for it, and that ModelRequest will not
-        # have the relevant instructions from the agent.
-
-        # While it's possible that you could have a message history where the most recent request has only tool returns,
-        # I believe there is no way to achieve that would _change_ the instructions without manually crafting the most
-        # recent message. That might make sense in principle for some usage pattern, but it's enough of an edge case
-        # that I think it's not worth worrying about, since you can work around this by inserting another ModelRequest
-        # with no parts at all immediately before the request that has the tool calls (that works because we only look
-        # at the two most recent ModelRequests here).
-
-        # If you have a use case where this causes pain, please open a GitHub issue and we can discuss alternatives.
-
-        if all(p.part_kind == 'tool-return' or p.part_kind == 'retry-prompt' for p in most_recent_request.parts):
-            return second_most_recent_request.instructions
-
-        return None
+        if instructions is None and len(last_two_requests) == 2:
+            most_recent_request = last_two_requests[0]
+            second_most_recent_request = last_two_requests[1]
+
+            # If we've gotten this far and the most recent request consists of only tool-return parts or retry-prompt parts,
+            # we use the instructions from the second-to-most-recent request. This is necessary because when handling
+            # result tools, we generate a "mock" ModelRequest with a tool-return part for it, and that ModelRequest will not
+            # have the relevant instructions from the agent.
+
+            # While it's possible that you could have a message history where the most recent request has only tool returns,
+            # I believe there is no way to achieve that would _change_ the instructions without manually crafting the most
+            # recent message. That might make sense in principle for some usage pattern, but it's enough of an edge case
+            # that I think it's not worth worrying about, since you can work around this by inserting another ModelRequest
+            # with no parts at all immediately before the request that has the tool calls (that works because we only look
+            # at the two most recent ModelRequests here).
+
+            # If you have a use case where this causes pain, please open a GitHub issue and we can discuss alternatives.
+
+            if all(p.part_kind == 'tool-return' or p.part_kind == 'retry-prompt' for p in most_recent_request.parts):
+                instructions = second_most_recent_request.instructions
+
+        if model_request_parameters and (output_instructions := model_request_parameters.prompted_output_instructions):
+            if instructions:
+                instructions = '\n\n'.join([instructions, output_instructions])
+            else:
+                instructions = output_instructions
+
+        return instructions
 
 
 @dataclass