pydantic-ai-slim 0.8.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

pydantic_ai/agent/abstract.py

@@ -2,12 +2,12 @@ from __future__ import annotations as _annotations
 
 import inspect
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterator, Mapping, Sequence
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from types import FrameType
-from typing import TYPE_CHECKING, Any, Callable, Generic, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, TypeAlias, cast, overload
 
-from typing_extensions import Self, TypeAlias, TypeIs, TypeVar
+from typing_extensions import Self, TypeIs, TypeVar
 
 from pydantic_graph import End
 from pydantic_graph._utils import get_event_loop
@@ -27,6 +27,7 @@ from ..run import AgentRun, AgentRunResult
 from ..settings import ModelSettings
 from ..tools import (
     AgentDepsT,
+    DeferredToolResults,
     RunContext,
     Tool,
     ToolFuncEither,
@@ -116,6 +117,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -133,6 +135,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -149,6 +152,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -180,6 +184,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             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.
@@ -201,6 +206,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             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,
@@ -225,6 +231,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -242,6 +249,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -258,6 +266,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -288,6 +297,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             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.
@@ -308,6 +318,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
                 user_prompt,
                 output_type=output_type,
                 message_history=message_history,
+                deferred_tool_results=deferred_tool_results,
                 model=model,
                 deps=deps,
                 model_settings=model_settings,
@@ -326,6 +337,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -343,6 +355,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -360,6 +373,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -398,6 +412,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             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.
@@ -424,6 +439,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             user_prompt,
             output_type=output_type,
             message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
             model=model,
             deps=deps,
             model_settings=model_settings,
@@ -436,8 +452,8 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             assert isinstance(first_node, _agent_graph.UserPromptNode)  # the first node should be a user prompt node
             node = first_node
             while True:
+                graph_ctx = agent_run.ctx
                 if self.is_model_request_node(node):
-                    graph_ctx = agent_run.ctx
                     async with node.stream(graph_ctx) as stream:
                         final_result_event = None
 
@@ -505,6 +521,17 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
                         await event_stream_handler(_agent_graph.build_run_context(agent_run.ctx), stream)
 
                 next_node = await agent_run.next(node)
+                if isinstance(next_node, End) and agent_run.result is not None:
+                    # A final output could have been produced by the CallToolsNode rather than the ModelRequestNode,
+                    # if a tool function raised CallDeferred or ApprovalRequired.
+                    # In this case there's no response to stream, but we still let the user access the output etc as normal.
+                    yield StreamedRunResult(
+                        graph_ctx.state.message_history,
+                        graph_ctx.deps.new_message_index,
+                        run_result=agent_run.result,
+                    )
+                    yielded = True
+                    break
                 if not isinstance(next_node, _agent_graph.AgentNode):
                     raise exceptions.AgentRunError(  # pragma: no cover
                         'Should have produced a StreamedRunResult before getting here'
@@ -521,6 +548,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -537,6 +565,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -554,6 +583,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
         *,
         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,
@@ -626,6 +656,7 @@ class AbstractAgent(Generic[AgentDepsT, OutputDataT], ABC):
             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.

pydantic_ai/agent/wrapper.py

@@ -15,6 +15,7 @@ from ..run import AgentRun
 from ..settings import ModelSettings
 from ..tools import (
     AgentDepsT,
+    DeferredToolResults,
     Tool,
     ToolFuncEither,
 )
@@ -72,6 +73,7 @@ class WrapperAgent(AbstractAgent[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,
@@ -88,6 +90,7 @@ class WrapperAgent(AbstractAgent[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,
@@ -104,6 +107,7 @@ class WrapperAgent(AbstractAgent[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,
@@ -176,6 +180,7 @@ class WrapperAgent(AbstractAgent[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.
@@ -191,6 +196,7 @@ class WrapperAgent(AbstractAgent[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/builtin_tools.py

@@ -9,7 +9,7 @@ from typing_extensions import TypedDict
 __all__ = ('AbstractBuiltinTool', 'WebSearchTool', 'WebSearchUserLocation', 'CodeExecutionTool', 'UrlContextTool')
 
 
-@dataclass
+@dataclass(kw_only=True)
 class AbstractBuiltinTool(ABC):
     """A builtin tool that can be used by an agent.
 
@@ -19,7 +19,7 @@ class AbstractBuiltinTool(ABC):
     """
 
 
-@dataclass
+@dataclass(kw_only=True)
 class WebSearchTool(AbstractBuiltinTool):
     """A builtin tool that allows your agent to search the web for information.
 

pydantic_ai/common_tools/duckduckgo.py

@@ -1,5 +1,5 @@
 import functools
-from dataclasses import dataclass
+from dataclasses import KW_ONLY, dataclass
 
 import anyio
 import anyio.to_thread
@@ -43,7 +43,9 @@ class DuckDuckGoSearchTool:
     client: DDGS
     """The DuckDuckGo search client."""
 
-    max_results: int | None = None
+    _: KW_ONLY
+
+    max_results: int | None
     """The maximum number of results. If None, returns results only from the first response."""
 
     async def __call__(self, query: str) -> list[DuckDuckGoResult]:

pydantic_ai/durable_exec/temporal/__init__.py

@@ -1,15 +1,24 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Sequence
+from collections.abc import AsyncIterator, Callable, Sequence
+from contextlib import AbstractAsyncContextManager
 from dataclasses import replace
-from typing import Any, Callable
+from typing import Any
 
 from pydantic.errors import PydanticUserError
-from temporalio.client import ClientConfig, Plugin as ClientPlugin
+from temporalio.client import ClientConfig, Plugin as ClientPlugin, WorkflowHistory
 from temporalio.contrib.pydantic import PydanticPayloadConverter, pydantic_data_converter
-from temporalio.converter import DefaultPayloadConverter
-from temporalio.worker import Plugin as WorkerPlugin, WorkerConfig
+from temporalio.converter import DataConverter, DefaultPayloadConverter
+from temporalio.service import ConnectConfig, ServiceClient
+from temporalio.worker import (
+    Plugin as WorkerPlugin,
+    Replayer,
+    ReplayerConfig,
+    Worker,
+    WorkerConfig,
+    WorkflowReplayResult,
+)
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 
 from ...exceptions import UserError
@@ -31,17 +40,15 @@ __all__ = [
 class PydanticAIPlugin(ClientPlugin, WorkerPlugin):
     """Temporal client and worker plugin for Pydantic AI."""
 
-    def configure_client(self, config: ClientConfig) -> ClientConfig:
-        if (data_converter := config.get('data_converter')) and data_converter.payload_converter_class not in (
-            DefaultPayloadConverter,
-            PydanticPayloadConverter,
-        ):
-            warnings.warn(  # pragma: no cover
-                'A non-default Temporal data converter was used which has been replaced with the Pydantic data converter.'
-            )
+    def init_client_plugin(self, next: ClientPlugin) -> None:
+        self.next_client_plugin = next
 
-        config['data_converter'] = pydantic_data_converter
-        return super().configure_client(config)
+    def init_worker_plugin(self, next: WorkerPlugin) -> None:
+        self.next_worker_plugin = next
+
+    def configure_client(self, config: ClientConfig) -> ClientConfig:
+        config['data_converter'] = self._get_new_data_converter(config.get('data_converter'))
+        return self.next_client_plugin.configure_client(config)
 
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         runner = config.get('workflow_runner')  # pyright: ignore[reportUnknownMemberType]
@@ -50,6 +57,8 @@ class PydanticAIPlugin(ClientPlugin, WorkerPlugin):
                 runner,
                 restrictions=runner.restrictions.with_passthrough_modules(
                     'pydantic_ai',
+                    'pydantic',
+                    'pydantic_core',
                     'logfire',
                     'rich',
                     'httpx',
@@ -67,7 +76,35 @@ class PydanticAIPlugin(ClientPlugin, WorkerPlugin):
             PydanticUserError,
         ]
 
-        return super().configure_worker(config)
+        return self.next_worker_plugin.configure_worker(config)
+
+    async def connect_service_client(self, config: ConnectConfig) -> ServiceClient:
+        return await self.next_client_plugin.connect_service_client(config)
+
+    async def run_worker(self, worker: Worker) -> None:
+        await self.next_worker_plugin.run_worker(worker)
+
+    def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:  # pragma: no cover
+        config['data_converter'] = self._get_new_data_converter(config.get('data_converter'))  # pyright: ignore[reportUnknownMemberType]
+        return self.next_worker_plugin.configure_replayer(config)
+
+    def run_replayer(
+        self,
+        replayer: Replayer,
+        histories: AsyncIterator[WorkflowHistory],
+    ) -> AbstractAsyncContextManager[AsyncIterator[WorkflowReplayResult]]:  # pragma: no cover
+        return self.next_worker_plugin.run_replayer(replayer, histories)
+
+    def _get_new_data_converter(self, converter: DataConverter | None) -> DataConverter:
+        if converter and converter.payload_converter_class not in (
+            DefaultPayloadConverter,
+            PydanticPayloadConverter,
+        ):
+            warnings.warn(  # pragma: no cover
+                'A non-default Temporal data converter was used which has been replaced with the Pydantic data converter.'
+            )
+
+        return pydantic_data_converter
 
 
 class AgentPlugin(WorkerPlugin):
@@ -76,8 +113,24 @@ class AgentPlugin(WorkerPlugin):
     def __init__(self, agent: TemporalAgent[Any, Any]):
         self.agent = agent
 
+    def init_worker_plugin(self, next: WorkerPlugin) -> None:
+        self.next_worker_plugin = next
+
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         activities: Sequence[Callable[..., Any]] = config.get('activities', [])  # pyright: ignore[reportUnknownMemberType]
         # Activities are checked for name conflicts by Temporal.
         config['activities'] = [*activities, *self.agent.temporal_activities]
-        return super().configure_worker(config)
+        return self.next_worker_plugin.configure_worker(config)
+
+    async def run_worker(self, worker: Worker) -> None:
+        await self.next_worker_plugin.run_worker(worker)
+
+    def configure_replayer(self, config: ReplayerConfig) -> ReplayerConfig:  # pragma: no cover
+        return self.next_worker_plugin.configure_replayer(config)
+
+    def run_replayer(
+        self,
+        replayer: Replayer,
+        histories: AsyncIterator[WorkflowHistory],
+    ) -> AbstractAsyncContextManager[AsyncIterator[WorkflowReplayResult]]:  # pragma: no cover
+        return self.next_worker_plugin.run_replayer(replayer, histories)