pydantic-ai-slim 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

pydantic_ai/agent/wrapper.py

@@ -0,0 +1,227 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import AsyncIterator, Iterator, Sequence
+from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
+from typing import Any, overload
+
+from .. import (
+    _utils,
+    messages as _messages,
+    models,
+    usage as _usage,
+)
+from ..output import OutputDataT, OutputSpec
+from ..run import AgentRun
+from ..settings import ModelSettings
+from ..tools import (
+    AgentDepsT,
+    Tool,
+    ToolFuncEither,
+)
+from ..toolsets import AbstractToolset
+from .abstract import AbstractAgent, EventStreamHandler, RunOutputDataT
+
+
+class WrapperAgent(AbstractAgent[AgentDepsT, OutputDataT]):
+    """Agent which wraps another agent.
+
+    Does nothing on its own, used as a base class.
+    """
+
+    def __init__(self, wrapped: AbstractAgent[AgentDepsT, OutputDataT]):
+        self.wrapped = wrapped
+
+    @property
+    def model(self) -> models.Model | models.KnownModelName | str | None:
+        return self.wrapped.model
+
+    @property
+    def name(self) -> str | None:
+        return self.wrapped.name
+
+    @name.setter
+    def name(self, value: str | None) -> None:
+        self.wrapped.name = value
+
+    @property
+    def deps_type(self) -> type:
+        return self.wrapped.deps_type
+
+    @property
+    def output_type(self) -> OutputSpec[OutputDataT]:
+        return self.wrapped.output_type
+
+    @property
+    def event_stream_handler(self) -> EventStreamHandler[AgentDepsT] | None:
+        return self.wrapped.event_stream_handler
+
+    @property
+    def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
+        return self.wrapped.toolsets
+
+    async def __aenter__(self) -> AbstractAgent[AgentDepsT, OutputDataT]:
+        return await self.wrapped.__aenter__()
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        return await self.wrapped.__aexit__(*args)
+
+    @overload
+    def iter(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: None = None,
+        message_history: list[_messages.ModelMessage] | 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.Usage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, OutputDataT]]: ...
+
+    @overload
+    def iter(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT],
+        message_history: list[_messages.ModelMessage] | 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.Usage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ...
+
+    @asynccontextmanager
+    async def iter(
+        self,
+        user_prompt: str | Sequence[_messages.UserContent] | None = None,
+        *,
+        output_type: OutputSpec[RunOutputDataT] | None = None,
+        message_history: list[_messages.ModelMessage] | 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.Usage | None = None,
+        infer_name: bool = True,
+        toolsets: Sequence[AbstractToolset[AgentDepsT]] | None = None,
+    ) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
+        """A contextmanager which can be used to iterate over the agent graph's nodes as they are executed.
+
+        This method builds an internal agent graph (using system prompts, tools and output schemas) and then returns an
+        `AgentRun` object. The `AgentRun` can be used to async-iterate over the nodes of the graph as they are
+        executed. This is the API to use if you want to consume the outputs coming from each LLM model response, or the
+        stream of events coming from the execution of tools.
+
+        The `AgentRun` also provides methods to access the full message history, new messages, and usage statistics,
+        and the final result of the run once it has completed.
+
+        For more details, see the documentation of `AgentRun`.
+
+        Example:
+        ```python
+        from pydantic_ai import Agent
+
+        agent = Agent('openai:gpt-4o')
+
+        async def main():
+            nodes = []
+            async with agent.iter('What is the capital of France?') as agent_run:
+                async for node in agent_run:
+                    nodes.append(node)
+            print(nodes)
+            '''
+            [
+                UserPromptNode(
+                    user_prompt='What is the capital of France?',
+                    instructions=None,
+                    instructions_functions=[],
+                    system_prompts=(),
+                    system_prompt_functions=[],
+                    system_prompt_dynamic_functions={},
+                ),
+                ModelRequestNode(
+                    request=ModelRequest(
+                        parts=[
+                            UserPromptPart(
+                                content='What is the capital of France?',
+                                timestamp=datetime.datetime(...),
+                            )
+                        ]
+                    )
+                ),
+                CallToolsNode(
+                    model_response=ModelResponse(
+                        parts=[TextPart(content='The capital of France is Paris.')],
+                        usage=Usage(
+                            requests=1, request_tokens=56, response_tokens=7, total_tokens=63
+                        ),
+                        model_name='gpt-4o',
+                        timestamp=datetime.datetime(...),
+                    )
+                ),
+                End(data=FinalResult(output='The capital of France is Paris.')),
+            ]
+            '''
+            print(agent_run.result.output)
+            #> The capital of France is Paris.
+        ```
+
+        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.
+            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:
+            The result of the run.
+        """
+        async with self.wrapped.iter(
+            user_prompt=user_prompt,
+            output_type=output_type,
+            message_history=message_history,
+            model=model,
+            deps=deps,
+            model_settings=model_settings,
+            usage_limits=usage_limits,
+            usage=usage,
+            infer_name=infer_name,
+            toolsets=toolsets,
+        ) as run:
+            yield run
+
+    @contextmanager
+    def override(
+        self,
+        *,
+        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,
+    ) -> Iterator[None]:
+        """Context manager to temporarily override agent dependencies, model, toolsets, or tools.
+
+        This is particularly useful when testing.
+        You can find an example of this [here](../testing.md#overriding-model-via-pytest-fixtures).
+
+        Args:
+            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.
+            tools: The tools to use instead of the tools registered with the agent.
+        """
+        with self.wrapped.override(deps=deps, model=model, toolsets=toolsets, tools=tools):
+            yield

pydantic_ai/direct.py

@@ -56,8 +56,8 @@ async def model_request(
         print(model_response)
         '''
         ModelResponse(
-            parts=[TextPart(content='Paris')],
-            usage=Usage(requests=1, request_tokens=56, response_tokens=1, total_tokens=57),
+            parts=[TextPart(content='The capital of France is Paris.')],
+            usage=Usage(requests=1, request_tokens=56, response_tokens=7, total_tokens=63),
             model_name='claude-3-5-haiku-latest',
             timestamp=datetime.datetime(...),
         )
@@ -109,8 +109,8 @@ def model_request_sync(
     print(model_response)
     '''
     ModelResponse(
-        parts=[TextPart(content='Paris')],
-        usage=Usage(requests=1, request_tokens=56, response_tokens=1, total_tokens=57),
+        parts=[TextPart(content='The capital of France is Paris.')],
+        usage=Usage(requests=1, request_tokens=56, response_tokens=7, total_tokens=63),
         model_name='claude-3-5-haiku-latest',
         timestamp=datetime.datetime(...),
     )
@@ -167,6 +167,7 @@ def model_request_stream(
             '''
             [
                 PartStartEvent(index=0, part=TextPart(content='Albert Einstein was ')),
+                FinalResultEvent(tool_name=None, tool_call_id=None),
                 PartDeltaEvent(
                     index=0, delta=TextPartDelta(content_delta='a German-born theoretical ')
                 ),
@@ -223,6 +224,7 @@ def model_request_stream_sync(
         '''
         [
             PartStartEvent(index=0, part=TextPart(content='Albert Einstein was ')),
+            FinalResultEvent(tool_name=None, tool_call_id=None),
             PartDeltaEvent(
                 index=0, delta=TextPartDelta(content_delta='a German-born theoretical ')
             ),
@@ -273,9 +275,7 @@ class StreamedResponseSync:
     """
 
     _async_stream_cm: AbstractAsyncContextManager[StreamedResponse]
-    _queue: queue.Queue[messages.ModelResponseStreamEvent | Exception | None] = field(
-        default_factory=queue.Queue, init=False
-    )
+    _queue: queue.Queue[messages.AgentStreamEvent | Exception | None] = field(default_factory=queue.Queue, init=False)
     _thread: threading.Thread | None = field(default=None, init=False)
     _stream_response: StreamedResponse | None = field(default=None, init=False)
     _exception: Exception | None = field(default=None, init=False)
@@ -295,8 +295,8 @@ class StreamedResponseSync:
     ) -> None:
         self._cleanup()
 
-    def __iter__(self) -> Iterator[messages.ModelResponseStreamEvent]:
-        """Stream the response as an iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
+    def __iter__(self) -> Iterator[messages.AgentStreamEvent]:
+        """Stream the response as an iterable of [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s."""
         self._check_context_manager_usage()
 
         while True:

pydantic_ai/durable_exec/temporal/__init__.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import warnings
+from collections.abc import Sequence
+from dataclasses import replace
+from typing import Any, Callable
+
+from pydantic.errors import PydanticUserError
+from temporalio.client import ClientConfig, Plugin as ClientPlugin
+from temporalio.contrib.pydantic import PydanticPayloadConverter, pydantic_data_converter
+from temporalio.converter import DefaultPayloadConverter
+from temporalio.worker import Plugin as WorkerPlugin, WorkerConfig
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
+
+from ...exceptions import UserError
+from ._agent import TemporalAgent
+from ._logfire import LogfirePlugin
+from ._run_context import TemporalRunContext
+from ._toolset import TemporalWrapperToolset
+
+__all__ = [
+    'TemporalAgent',
+    'PydanticAIPlugin',
+    'LogfirePlugin',
+    'AgentPlugin',
+    'TemporalRunContext',
+    'TemporalWrapperToolset',
+]
+
+
+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.'
+            )
+
+        config['data_converter'] = pydantic_data_converter
+        return super().configure_client(config)
+
+    def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
+        runner = config.get('workflow_runner')  # pyright: ignore[reportUnknownMemberType]
+        if isinstance(runner, SandboxedWorkflowRunner):  # pragma: no branch
+            config['workflow_runner'] = replace(
+                runner,
+                restrictions=runner.restrictions.with_passthrough_modules(
+                    'pydantic_ai',
+                    'logfire',
+                    'rich',
+                    'httpx',
+                    # Imported inside `logfire._internal.json_encoder` when running `logfire.info` inside an activity with attributes to serialize
+                    'attrs',
+                    # Imported inside `logfire._internal.json_schema` when running `logfire.info` inside an activity with attributes to serialize
+                    'numpy',
+                    'pandas',
+                ),
+            )
+
+        config['workflow_failure_exception_types'] = [
+            *config.get('workflow_failure_exception_types', []),  # pyright: ignore[reportUnknownMemberType]
+            UserError,
+            PydanticUserError,
+        ]
+
+        return super().configure_worker(config)
+
+
+class AgentPlugin(WorkerPlugin):
+    """Temporal worker plugin for a specific Pydantic AI agent."""
+
+    def __init__(self, agent: TemporalAgent[Any, Any]):
+        self.agent = agent
+
+    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)