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

pydantic_ai/durable_exec/dbos/_mcp_server.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from abc import ABC
+from collections.abc import Callable
+from typing import Any
+
+from dbos import DBOS
+from typing_extensions import Self
+
+from pydantic_ai.mcp import MCPServer, ToolResult
+from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
+from pydantic_ai.toolsets.wrapper import WrapperToolset
+
+from ._utils import StepConfig
+
+
+class DBOSMCPServer(WrapperToolset[AgentDepsT], ABC):
+    """A wrapper for MCPServer that integrates with DBOS, turning call_tool and get_tools to DBOS steps."""
+
+    def __init__(
+        self,
+        wrapped: MCPServer,
+        *,
+        step_name_prefix: str,
+        step_config: StepConfig,
+    ):
+        super().__init__(wrapped)
+        self._step_config = step_config or {}
+        self._step_name_prefix = step_name_prefix
+        id_suffix = f'__{wrapped.id}' if wrapped.id else ''
+        self._name = f'{step_name_prefix}__mcp_server{id_suffix}'
+
+        # Wrap get_tools in a DBOS step.
+        @DBOS.step(
+            name=f'{self._name}.get_tools',
+            **self._step_config,
+        )
+        async def wrapped_get_tools_step(
+            ctx: RunContext[AgentDepsT],
+        ) -> dict[str, ToolsetTool[AgentDepsT]]:
+            return await super(DBOSMCPServer, self).get_tools(ctx)
+
+        self._dbos_wrapped_get_tools_step = wrapped_get_tools_step
+
+        # Wrap call_tool in a DBOS step.
+        @DBOS.step(
+            name=f'{self._name}.call_tool',
+            **self._step_config,
+        )
+        async def wrapped_call_tool_step(
+            name: str,
+            tool_args: dict[str, Any],
+            ctx: RunContext[AgentDepsT],
+            tool: ToolsetTool[AgentDepsT],
+        ) -> ToolResult:
+            return await super(DBOSMCPServer, self).call_tool(name, tool_args, ctx, tool)
+
+        self._dbos_wrapped_call_tool_step = wrapped_call_tool_step
+
+    @property
+    def id(self) -> str | None:
+        return self.wrapped.id
+
+    async def __aenter__(self) -> Self:
+        # The wrapped MCPServer enters itself around listing and calling tools
+        # so we don't need to enter it here (nor could we because we're not inside a DBOS step).
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        return None
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        # DBOS-ified toolsets cannot be swapped out after the fact.
+        return self
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        return await self._dbos_wrapped_get_tools_step(ctx)
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[AgentDepsT],
+        tool: ToolsetTool[AgentDepsT],
+    ) -> ToolResult:
+        return await self._dbos_wrapped_call_tool_step(name, tool_args, ctx, tool)

pydantic_ai/durable_exec/dbos/_model.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from typing import Any
+
+from dbos import DBOS
+
+from pydantic_ai.agent import EventStreamHandler
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelResponse,
+    ModelResponseStreamEvent,
+)
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
+from pydantic_ai.models.wrapper import WrapperModel
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.tools import RunContext
+from pydantic_ai.usage import RequestUsage
+
+from ._utils import StepConfig
+
+
+class DBOSStreamedResponse(StreamedResponse):
+    def __init__(self, model_request_parameters: ModelRequestParameters, response: ModelResponse):
+        super().__init__(model_request_parameters)
+        self.response = response
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        return
+        # noinspection PyUnreachableCode
+        yield
+
+    def get(self) -> ModelResponse:
+        return self.response
+
+    def usage(self) -> RequestUsage:
+        return self.response.usage  # pragma: no cover
+
+    @property
+    def model_name(self) -> str:
+        return self.response.model_name or ''  # pragma: no cover
+
+    @property
+    def provider_name(self) -> str:
+        return self.response.provider_name or ''  # pragma: no cover
+
+    @property
+    def timestamp(self) -> datetime:
+        return self.response.timestamp  # pragma: no cover
+
+
+class DBOSModel(WrapperModel):
+    """A wrapper for Model that integrates with DBOS, turning request and request_stream to DBOS steps."""
+
+    def __init__(
+        self,
+        model: Model,
+        *,
+        step_name_prefix: str,
+        step_config: StepConfig,
+        event_stream_handler: EventStreamHandler[Any] | None = None,
+    ):
+        super().__init__(model)
+        self.step_config = step_config
+        self.event_stream_handler = event_stream_handler
+        self._step_name_prefix = step_name_prefix
+
+        # Wrap the request in a DBOS step.
+        @DBOS.step(
+            name=f'{self._step_name_prefix}__model.request',
+            **self.step_config,
+        )
+        async def wrapped_request_step(
+            messages: list[ModelMessage],
+            model_settings: ModelSettings | None,
+            model_request_parameters: ModelRequestParameters,
+        ) -> ModelResponse:
+            return await super(DBOSModel, self).request(messages, model_settings, model_request_parameters)
+
+        self._dbos_wrapped_request_step = wrapped_request_step
+
+        # Wrap the request_stream in a DBOS step.
+        @DBOS.step(
+            name=f'{self._step_name_prefix}__model.request_stream',
+            **self.step_config,
+        )
+        async def wrapped_request_stream_step(
+            messages: list[ModelMessage],
+            model_settings: ModelSettings | None,
+            model_request_parameters: ModelRequestParameters,
+            run_context: RunContext[Any] | None = None,
+        ) -> ModelResponse:
+            async with super(DBOSModel, self).request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as streamed_response:
+                if self.event_stream_handler is not None:
+                    assert run_context is not None, (
+                        'A DBOS model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+                    )
+                    await self.event_stream_handler(run_context, streamed_response)
+
+                async for _ in streamed_response:
+                    pass
+            return streamed_response.get()
+
+        self._dbos_wrapped_request_stream_step = wrapped_request_stream_step
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        return await self._dbos_wrapped_request_step(messages, model_settings, model_request_parameters)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        # If not in a workflow (could be in a step), just call the wrapped request_stream method.
+        if DBOS.workflow_id is None or DBOS.step_id is not None:
+            async with super().request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as streamed_response:
+                yield streamed_response
+                return
+
+        response = await self._dbos_wrapped_request_stream_step(
+            messages, model_settings, model_request_parameters, run_context
+        )
+        yield DBOSStreamedResponse(model_request_parameters, response)

pydantic_ai/durable_exec/dbos/_utils.py

@@ -0,0 +1,10 @@
+from typing_extensions import TypedDict
+
+
+class StepConfig(TypedDict, total=False):
+    """Configuration for a step in the DBOS workflow."""
+
+    retries_allowed: bool
+    interval_seconds: float
+    max_attempts: int
+    backoff_rate: float

pydantic_ai/durable_exec/temporal/_agent.py

@@ -1,14 +1,16 @@
 from __future__ import annotations
 
-from collections.abc import AsyncIterator, Callable, Iterator, Sequence
+from collections.abc import AsyncIterable, AsyncIterator, Callable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from contextvars import ContextVar
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any, Literal, overload
 
+from pydantic import ConfigDict, with_config
 from pydantic.errors import PydanticUserError
 from pydantic_core import PydanticSerializationError
-from temporalio import workflow
+from temporalio import activity, workflow
 from temporalio.common import RetryPolicy
 from temporalio.workflow import ActivityConfig
 from typing_extensions import Never
@@ -21,7 +23,6 @@ from pydantic_ai import (
 )
 from pydantic_ai._run_context import AgentDepsT
 from pydantic_ai.agent import AbstractAgent, AgentRun, AgentRunResult, EventStreamHandler, RunOutputDataT, WrapperAgent
-from pydantic_ai.durable_exec.temporal._run_context import TemporalRunContext
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.models import Model
 from pydantic_ai.output import OutputDataT, OutputSpec
@@ -29,15 +30,24 @@ from pydantic_ai.result import StreamedRunResult
 from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import (
     DeferredToolResults,
+    RunContext,
     Tool,
     ToolFuncEither,
 )
 from pydantic_ai.toolsets import AbstractToolset
 
 from ._model import TemporalModel
+from ._run_context import TemporalRunContext
 from ._toolset import TemporalWrapperToolset, temporalize_toolset
 
 
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class _EventStreamHandlerParams:
+    event: _messages.AgentStreamEvent
+    serialized_run_context: Any
+
+
 class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
     def __init__(
         self,
@@ -86,6 +96,10 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
         """
         super().__init__(wrapped)
 
+        self._name = name
+        self._event_stream_handler = event_stream_handler
+        self.run_context_type = run_context_type
+
         # start_to_close_timeout is required
         activity_config = activity_config or ActivityConfig(start_to_close_timeout=timedelta(seconds=60))
 
@@ -97,13 +111,13 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
             PydanticUserError.__name__,
         ]
         activity_config['retry_policy'] = retry_policy
+        self.activity_config = activity_config
 
         model_activity_config = model_activity_config or {}
         toolset_activity_config = toolset_activity_config or {}
         tool_activity_config = tool_activity_config or {}
 
-        self._name = name or wrapped.name
-        if self._name is None:
+        if self.name is None:
             raise UserError(
                 "An agent needs to have a unique `name` in order to be used with Temporal. The name will be used to identify the agent's activities within the workflow."
             )
@@ -116,13 +130,33 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
                 'An agent needs to have a `model` in order to be used with Temporal, it cannot be set at agent run time.'
             )
 
+        async def event_stream_handler_activity(params: _EventStreamHandlerParams, deps: AgentDepsT) -> None:
+            # We can never get here without an `event_stream_handler`, as `TemporalAgent.run_stream` and `TemporalAgent.iter` raise an error saying to use `TemporalAgent.run` instead,
+            # and that only ends up calling `event_stream_handler` if it is set.
+            assert self.event_stream_handler is not None
+
+            run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
+
+            async def streamed_response():
+                yield params.event
+
+            await self.event_stream_handler(run_context, streamed_response())
+
+        # Set type hint explicitly so that Temporal can take care of serialization and deserialization
+        event_stream_handler_activity.__annotations__['deps'] = self.deps_type
+
+        self.event_stream_handler_activity = activity.defn(name=f'{activity_name_prefix}__event_stream_handler')(
+            event_stream_handler_activity
+        )
+        activities.append(self.event_stream_handler_activity)
+
         temporal_model = TemporalModel(
             wrapped.model,
             activity_name_prefix=activity_name_prefix,
             activity_config=activity_config | model_activity_config,
             deps_type=self.deps_type,
-            run_context_type=run_context_type,
-            event_stream_handler=event_stream_handler or wrapped.event_stream_handler,
+            run_context_type=self.run_context_type,
+            event_stream_handler=self.event_stream_handler,
         )
         activities.extend(temporal_model.temporal_activities)
 
@@ -139,7 +173,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
                 activity_config | toolset_activity_config.get(id, {}),
                 tool_activity_config.get(id, {}),
                 self.deps_type,
-                run_context_type,
+                self.run_context_type,
             )
             if isinstance(toolset, TemporalWrapperToolset):
                 activities.extend(toolset.temporal_activities)
@@ -155,7 +189,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
 
     @property
     def name(self) -> str | None:
-        return self._name
+        return self._name or super().name
 
     @name.setter
     def name(self, value: str | None) -> None:  # pragma: no cover
@@ -167,6 +201,33 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
     def model(self) -> Model:
         return self._model
 
+    @property
+    def event_stream_handler(self) -> EventStreamHandler[AgentDepsT] | None:
+        handler = self._event_stream_handler or super().event_stream_handler
+        if handler is None:
+            return None
+        elif workflow.in_workflow():
+            return self._call_event_stream_handler_activity
+        else:
+            return handler
+
+    async def _call_event_stream_handler_activity(
+        self, ctx: RunContext[AgentDepsT], stream: AsyncIterable[_messages.AgentStreamEvent]
+    ) -> None:
+        serialized_run_context = self.run_context_type.serialize_run_context(ctx)
+        async for event in stream:
+            await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+                activity=self.event_stream_handler_activity,
+                args=[
+                    _EventStreamHandlerParams(
+                        event=event,
+                        serialized_run_context=serialized_run_context,
+                    ),
+                    ctx.deps,
+                ],
+                **self.activity_config,
+            )
+
     @property
     def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
         with self._temporal_overrides():
@@ -296,7 +357,7 @@ class TemporalAgent(WrapperAgent[AgentDepsT, OutputDataT]):
                 usage=usage,
                 infer_name=infer_name,
                 toolsets=toolsets,
-                event_stream_handler=event_stream_handler,
+                event_stream_handler=event_stream_handler or self.event_stream_handler,
                 **_deprecated_kwargs,
             )
 

pydantic_ai/exceptions.py

@@ -65,7 +65,7 @@ class ModelRetry(Exception):
 class CallDeferred(Exception):
     """Exception to raise when a tool call should be deferred.
 
-    See [tools docs](../tools.md#deferred-tools) for more information.
+    See [tools docs](../deferred-tools.md#deferred-tools) for more information.
     """
 
     pass
@@ -74,7 +74,7 @@ class CallDeferred(Exception):
 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.
+    See [tools docs](../deferred-tools.md#human-in-the-loop-tool-approval) for more information.
     """
 
     pass

pydantic_ai/mcp.py

@@ -20,6 +20,7 @@ from typing_extensions import Self, assert_never, deprecated
 
 from pydantic_ai.tools import RunContext, ToolDefinition
 
+from .direct import model_request
 from .toolsets.abstract import AbstractToolset, ToolsetTool
 
 try:
@@ -300,6 +301,8 @@ class MCPServer(AbstractToolset[Any], ABC):
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
+        if self._running_count == 0:
+            raise ValueError('MCPServer.__aexit__ called more times than __aenter__')
         async with self._enter_lock:
             self._running_count -= 1
             if self._running_count == 0 and self._exit_stack is not None:
@@ -327,11 +330,7 @@ class MCPServer(AbstractToolset[Any], ABC):
         if stop_sequences := params.stopSequences:  # pragma: no branch
             model_settings['stop_sequences'] = stop_sequences
 
-        model_response = await self.sampling_model.request(
-            pai_messages,
-            model_settings,
-            models.ModelRequestParameters(),
-        )
+        model_response = await model_request(self.sampling_model, pai_messages, model_settings=model_settings)
         return mcp_types.CreateMessageResult(
             role='assistant',
             content=_mcp.map_from_model_response(model_response),
@@ -401,16 +400,7 @@ class MCPServerStdio(MCPServer):
     from pydantic_ai.mcp import MCPServerStdio
 
     server = MCPServerStdio(  # (1)!
-        'deno',
-        args=[
-            'run',
-            '-N',
-            '-R=node_modules',
-            '-W=node_modules',
-            '--node-modules-dir=auto',
-            'jsr:@pydantic/mcp-run-python',
-            'stdio',
-        ]
+        'uv', args=['run', 'mcp-run-python', 'stdio'], timeout=10
     )
     agent = Agent('openai:gpt-4o', toolsets=[server])
 
@@ -419,7 +409,7 @@ class MCPServerStdio(MCPServer):
             ...
     ```
 
-    1. See [MCP Run Python](../mcp/run-python.md) for more information.
+    1. See [MCP Run Python](https://github.com/pydantic/mcp-run-python) for more information.
     2. This will start the server as a subprocess and connect to it.
     """
 
@@ -455,6 +445,7 @@ class MCPServerStdio(MCPServer):
         self,
         command: str,
         args: Sequence[str],
+        *,
         env: dict[str, str] | None = None,
         cwd: str | Path | None = None,
         tool_prefix: str | None = None,
@@ -467,7 +458,6 @@ class MCPServerStdio(MCPServer):
         sampling_model: models.Model | None = None,
         max_retries: int = 1,
         elicitation_callback: ElicitationFnT | None = None,
-        *,
         id: str | None = None,
     ):
         """Build a new MCP server.
@@ -527,7 +517,7 @@ class MCPServerStdio(MCPServer):
             f'args={self.args!r}',
         ]
         if self.id:
-            repr_args.append(f'id={self.id!r}')  # pragma: no cover
+            repr_args.append(f'id={self.id!r}')
         return f'{self.__class__.__name__}({", ".join(repr_args)})'
 
 
@@ -581,8 +571,8 @@ class _MCPServerHTTP(MCPServer):
 
     def __init__(
         self,
-        *,
         url: str,
+        *,
         headers: dict[str, str] | None = None,
         http_client: httpx.AsyncClient | None = None,
         id: str | None = None,
@@ -732,16 +722,15 @@ class MCPServerSSE(_MCPServerHTTP):
     from pydantic_ai import Agent
     from pydantic_ai.mcp import MCPServerSSE
 
-    server = MCPServerSSE('http://localhost:3001/sse')  # (1)!
+    server = MCPServerSSE('http://localhost:3001/sse')
     agent = Agent('openai:gpt-4o', toolsets=[server])
 
     async def main():
-        async with agent:  # (2)!
+        async with agent:  # (1)!
             ...
     ```
 
-    1. E.g. you might be connecting to a server run with [`mcp-run-python`](../mcp/run-python.md).
-    2. This will connect to a server running on `localhost:3001`.
+    1. This will connect to a server running on `localhost:3001`.
     """
 
     @property
@@ -765,7 +754,7 @@ class MCPServerHTTP(MCPServerSSE):
     from pydantic_ai import Agent
     from pydantic_ai.mcp import MCPServerHTTP
 
-    server = MCPServerHTTP('http://localhost:3001/sse')  # (1)!
+    server = MCPServerHTTP('http://localhost:3001/sse')
     agent = Agent('openai:gpt-4o', toolsets=[server])
 
     async def main():
@@ -773,8 +762,7 @@ class MCPServerHTTP(MCPServerSSE):
             ...
     ```
 
-    1. E.g. you might be connecting to a server run with [`mcp-run-python`](../mcp/run-python.md).
-    2. This will connect to a server running on `localhost:3001`.
+    1. This will connect to a server running on `localhost:3001`.
     """