pydantic-ai-slim 0.8.1__py3-none-any.whl → 1.0.0b1__py3-none-any.whl

pydantic_ai/durable_exec/temporal/_agent.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Iterator, Sequence
+from collections.abc import AsyncIterator, Callable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from datetime import timedelta
-from typing import Any, Callable, Literal, overload
+from typing import Any, Literal, overload
 
 from pydantic.errors import PydanticUserError
 from pydantic_core import PydanticSerializationError
@@ -28,6 +28,7 @@ from pydantic_ai.output import OutputDataT, OutputSpec
 from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import (
+    DeferredToolResults,
     Tool,
     ToolFuncEither,
 )
@@ -196,6 +197,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -213,6 +215,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -229,6 +232,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -261,6 +265,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             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.
@@ -283,6 +288,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
                 user_prompt,
                 output_type=output_type,
                 message_history=message_history,
+                deferred_tool_results=deferred_tool_results,
                 model=model,
                 deps=deps,
                 model_settings=model_settings,
@@ -301,6 +307,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -318,6 +325,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -334,6 +342,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -365,6 +374,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             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.
@@ -386,6 +396,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             user_prompt,
             output_type=output_type,
             message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
             model=model,
             deps=deps,
             model_settings=model_settings,
@@ -404,6 +415,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -421,6 +433,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -438,6 +451,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -467,6 +481,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             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.
@@ -490,6 +505,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             user_prompt,
             output_type=output_type,
             message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
             model=model,
             deps=deps,
             model_settings=model_settings,
@@ -509,6 +525,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -526,6 +543,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -543,6 +561,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         *,
         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,
@@ -616,6 +635,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             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.
@@ -648,6 +668,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             user_prompt=user_prompt,
             output_type=output_type,
             message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
             model=model,
             deps=deps,
             model_settings=model_settings,

pydantic_ai/durable_exec/temporal/_function_toolset.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, Literal
+from typing import Annotated, Any, Literal, assert_never
 
-from pydantic import ConfigDict, with_config
+from pydantic import ConfigDict, Discriminator, with_config
 from temporalio import activity, workflow
 from temporalio.workflow import ActivityConfig
 
-from pydantic_ai.exceptions import UserError
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
 from pydantic_ai.tools import AgentDepsT, RunContext
 from pydantic_ai.toolsets import FunctionToolset, ToolsetTool
 from pydantic_ai.toolsets.function import FunctionToolsetTool
@@ -24,6 +25,34 @@ class _CallToolParams:
     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'),
+]
+
+
 class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
     def __init__(
         self,
@@ -40,7 +69,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) -> Any:
+        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:
@@ -54,7 +83,15 @@ 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)
-            return await self.wrapped.call_tool(name, args_dict, ctx, tool)
+            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)
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         call_tool_activity.__annotations__['deps'] = deps_type
@@ -85,7 +122,7 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
 
         tool_activity_config = self.activity_config | tool_activity_config
         serialized_run_context = self.run_context_type.serialize_run_context(ctx)
-        return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+        result = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.call_tool_activity,
             args=[
                 _CallToolParams(
@@ -97,3 +134,13 @@ class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
             ],
             **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/_logfire.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Callable
+from collections.abc import Callable
 
 from logfire import Logfire
 from opentelemetry.trace import get_tracer

pydantic_ai/durable_exec/temporal/_mcp_server.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, Literal
+from typing import Any, Literal
 
 from pydantic import ConfigDict, with_config
 from temporalio import activity, workflow

pydantic_ai/durable_exec/temporal/_model.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Callable
+from typing import Any
 
 from pydantic import ConfigDict, with_config
 from temporalio import activity, workflow

pydantic_ai/durable_exec/temporal/_run_context.py

@@ -9,7 +9,7 @@ from pydantic_ai.tools import AgentDepsT, 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`, `retry` and `run_step` attributes will be available.
+    By default, only the `deps`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry` and `run_step` 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].
     """
 
@@ -40,6 +40,7 @@ class TemporalRunContext(RunContext[AgentDepsT]):
             'retries': ctx.retries,
             'tool_call_id': ctx.tool_call_id,
             'tool_name': ctx.tool_name,
+            'tool_call_approved': ctx.tool_call_approved,
             'retry': ctx.retry,
             'run_step': ctx.run_step,
         }

pydantic_ai/durable_exec/temporal/_toolset.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Literal
+from collections.abc import Callable
+from typing import Any, Literal
 
 from temporalio.workflow import ActivityConfig
 

pydantic_ai/exceptions.py

@@ -2,7 +2,9 @@ from __future__ import annotations as _annotations
 
 import json
 import sys
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+from pydantic_core import core_schema
 
 if sys.version_info < (3, 11):
     from exceptiongroup import ExceptionGroup as ExceptionGroup  # pragma: lax no cover
@@ -14,6 +16,8 @@ if TYPE_CHECKING:
 
 __all__ = (
     'ModelRetry',
+    'CallDeferred',
+    'ApprovalRequired',
     'UserError',
     'AgentRunError',
     'UnexpectedModelBehavior',
@@ -24,7 +28,7 @@ __all__ = (
 
 
 class ModelRetry(Exception):
-    """Exception raised when a tool function should be retried.
+    """Exception to raise when a tool function should be retried.
 
     The agent will return the message to the model and ask it to try calling the function/tool again.
     """
@@ -36,6 +40,45 @@ class ModelRetry(Exception):
         self.message = message
         super().__init__(message)
 
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, self.__class__) and other.message == self.message
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _: Any, __: Any) -> core_schema.CoreSchema:
+        """Pydantic core schema to allow `ModelRetry` to be (de)serialized."""
+        schema = core_schema.typed_dict_schema(
+            {
+                'message': core_schema.typed_dict_field(core_schema.str_schema()),
+                'kind': core_schema.typed_dict_field(core_schema.literal_schema(['model-retry'])),
+            }
+        )
+        return core_schema.no_info_after_validator_function(
+            lambda dct: ModelRetry(dct['message']),
+            schema,
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda x: {'message': x.message, 'kind': 'model-retry'},
+                return_schema=schema,
+            ),
+        )
+
+
+class CallDeferred(Exception):
+    """Exception to raise when a tool call should be deferred.
+
+    See [tools docs](../tools.md#deferred-tools) for more information.
+    """
+
+    pass
+
+
+class ApprovalRequired(Exception):
+    """Exception to raise when a tool call requires human-in-the-loop approval.
+
+    See [tools docs](../tools.md#human-in-the-loop-tool-approval) for more information.
+    """
+
+    pass
+
 
 class UserError(RuntimeError):
     """Error caused by a usage mistake by the application developer — You!"""

pydantic_ai/format_prompt.py

@@ -72,9 +72,9 @@ class _ToXml:
             element.text = self.none_str
         elif isinstance(value, str):
             element.text = value
-        elif isinstance(value, (bytes, bytearray)):
+        elif isinstance(value, bytes | bytearray):
             element.text = value.decode(errors='ignore')
-        elif isinstance(value, (bool, int, float)):
+        elif isinstance(value, bool | int | float):
             element.text = str(value)
         elif isinstance(value, date):
             element.text = value.isoformat()

pydantic_ai/mcp.py

@@ -5,12 +5,12 @@ import functools
 import warnings
 from abc import ABC, abstractmethod
 from asyncio import Lock
-from collections.abc import AsyncIterator, Awaitable, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
 from dataclasses import field, replace
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any
 
 import anyio
 import httpx