WIP: temporalize_agent

Author: DouweM

pydantic_ai_slim/pydantic_ai/agent.py

@@ -1717,6 +1717,14 @@ def _get_toolset(
 
         return CombinedToolset(all_toolsets)
 
+    @property
+    def toolset(self) -> AbstractToolset[AgentDepsT]:
+        """The complete toolset that will be available to the model during an agent run.
+
+        This will include function tools registered directly to the agent, output tools, and user-provided toolsets including MCP servers.
+        """
+        return self._get_toolset()
+
     def _infer_name(self, function_frame: FrameType | None) -> None:
         """Infer the agent name from the call frame.
 

pydantic_ai_slim/pydantic_ai/temporal/__init__.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any, Callable
+
+from temporalio.common import Priority, RetryPolicy
+from temporalio.workflow import ActivityCancellationType, VersioningIntent
+
+from pydantic_ai._run_context import AgentDepsT, RunContext
+
+
+class _TemporalRunContext(RunContext[AgentDepsT]):
+    _data: dict[str, Any]
+
+    def __init__(self, **kwargs: Any):
+        self._data = kwargs
+        setattr(
+            self,
+            '__dataclass_fields__',
+            {name: field for name, field in RunContext.__dataclass_fields__.items() if name in kwargs},
+        )
+
+    def __getattribute__(self, name: str) -> Any:
+        try:
+            return super().__getattribute__(name)
+        except AttributeError as e:
+            data = super().__getattribute__('_data')
+            if name in data:
+                return data[name]
+            raise e  # TODO: Explain how to make a new run context attribute available
+
+    @classmethod
+    def serialize_run_context(cls, ctx: RunContext[AgentDepsT]) -> dict[str, Any]:
+        return {
+            'deps': ctx.deps,
+            'retries': ctx.retries,
+            'tool_call_id': ctx.tool_call_id,
+            'tool_name': ctx.tool_name,
+            'retry': ctx.retry,
+            'run_step': ctx.run_step,
+        }
+
+    @classmethod
+    def deserialize_run_context(cls, ctx: dict[str, Any]) -> RunContext[AgentDepsT]:
+        return cls(**ctx)
+
+
+@dataclass
+class TemporalSettings:
+    """Settings for Temporal `execute_activity` and Pydantic AI-specific Temporal activity behavior."""
+
+    # Temporal settings
+    task_queue: str | None = None
+    schedule_to_close_timeout: timedelta | None = None
+    schedule_to_start_timeout: timedelta | None = None
+    start_to_close_timeout: timedelta | None = None
+    heartbeat_timeout: timedelta | None = None
+    retry_policy: RetryPolicy | None = None
+    cancellation_type: ActivityCancellationType = ActivityCancellationType.TRY_CANCEL
+    activity_id: str | None = None
+    versioning_intent: VersioningIntent | None = None
+    summary: str | None = None
+    priority: Priority = Priority.default
+
+    # Pydantic AI specific
+    tool_settings: dict[str, dict[str, TemporalSettings]] | None = None
+
+    def for_tool(self, toolset_id: str, tool_id: str) -> TemporalSettings:
+        if self.tool_settings is None:
+            return self
+        return self.tool_settings.get(toolset_id, {}).get(tool_id, self)
+
+    serialize_run_context: Callable[[RunContext], Any] = _TemporalRunContext.serialize_run_context
+    deserialize_run_context: Callable[[dict[str, Any]], RunContext] = _TemporalRunContext.deserialize_run_context
+
+    @property
+    def execute_activity_kwargs(self) -> dict[str, Any]:
+        return {
+            'task_queue': self.task_queue,
+            'schedule_to_close_timeout': self.schedule_to_close_timeout,
+            'schedule_to_start_timeout': self.schedule_to_start_timeout,
+            'start_to_close_timeout': self.start_to_close_timeout,
+            'heartbeat_timeout': self.heartbeat_timeout,
+            'retry_policy': self.retry_policy,
+            'cancellation_type': self.cancellation_type,
+            'activity_id': self.activity_id,
+            'versioning_intent': self.versioning_intent,
+            'summary': self.summary,
+            'priority': self.priority,
+        }
+
+
+def initialize_temporal():
+    """Explicitly import types without which Temporal will not be able to serialize/deserialize `ModelMessage`s."""
+    from pydantic_ai.messages import (  # noqa F401
+        ModelResponse,  # pyright: ignore[reportUnusedImport]
+        ImageUrl,  # pyright: ignore[reportUnusedImport]
+        AudioUrl,  # pyright: ignore[reportUnusedImport]
+        DocumentUrl,  # pyright: ignore[reportUnusedImport]
+        VideoUrl,  # pyright: ignore[reportUnusedImport]
+        BinaryContent,  # pyright: ignore[reportUnusedImport]
+        UserContent,  # pyright: ignore[reportUnusedImport]
+    )

pydantic_ai_slim/pydantic_ai/temporal/agent.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from typing import Any, Callable, TypeVar
+
+from pydantic_ai.agent import Agent
+from pydantic_ai.mcp import MCPServer
+from pydantic_ai.toolsets.abstract import AbstractToolset
+from pydantic_ai.toolsets.function import FunctionToolset
+
+from ..models import Model
+from . import TemporalSettings
+from .function_toolset import temporalize_function_toolset
+from .mcp_server import temporalize_mcp_server
+from .model import temporalize_model
+
+T = TypeVar('T', bound=AbstractToolset, covariant=True)
+
+
+def temporalize_toolset(toolset: AbstractToolset, settings: TemporalSettings | None) -> list[Callable[..., Any]]:
+    """Temporalize a toolset.
+
+    Args:
+        toolset: The toolset to temporalize.
+        settings: The temporal settings to use.
+    """
+    if isinstance(toolset, FunctionToolset):
+        return temporalize_function_toolset(toolset, settings)
+    elif isinstance(toolset, MCPServer):
+        return temporalize_mcp_server(toolset, settings)
+    else:
+        return []
+
+
+def temporalize_agent(
+    agent: Agent,
+    settings: TemporalSettings | None = None,
+    temporalize_toolset_func: Callable[
+        [AbstractToolset, TemporalSettings | None], list[Callable[..., Any]]
+    ] = temporalize_toolset,
+) -> list[Callable[..., Any]]:
+    """Temporalize an agent.
+
+    Args:
+        agent: The agent to temporalize.
+        settings: The temporal settings to use.
+        temporalize_toolset_func: The function to use to temporalize the toolsets.
+    """
+    if existing_activities := getattr(agent, '__temporal_activities', None):
+        return existing_activities
+
+    settings = settings or TemporalSettings()
+
+    # TODO: Doesn't consider model/toolsets passed at iter time.
+
+    activities: list[Callable[..., Any]] = []
+    if isinstance(agent.model, Model):
+        activities.extend(temporalize_model(agent.model, settings))
+
+    def temporalize_toolset(toolset: AbstractToolset) -> None:
+        activities.extend(temporalize_toolset_func(toolset, settings))
+
+    agent.toolset.apply(temporalize_toolset)
+
+    setattr(agent, '__temporal_activities', activities)
+    return activities

pydantic_ai_slim/pydantic_ai/temporal/function_toolset.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from pydantic import ConfigDict, with_config
+from temporalio import activity, workflow
+
+from pydantic_ai.toolsets.function import FunctionToolset
+
+from .._run_context import RunContext
+from ..toolsets import ToolsetTool
+from . import TemporalSettings
+
+
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class _CallToolParams:
+    name: str
+    tool_args: dict[str, Any]
+    serialized_run_context: Any
+
+
+def temporalize_function_toolset(
+    toolset: FunctionToolset,
+    settings: TemporalSettings | None = None,
+) -> list[Callable[..., Any]]:
+    """Temporalize a function toolset.
+
+    Args:
+        toolset: The function toolset to temporalize.
+        settings: The temporal settings to use.
+    """
+    if activities := getattr(toolset, '__temporal_activities', None):
+        return activities
+
+    id = toolset.id
+    if not id:
+        raise ValueError(
+            "A function toolset needs to have an ID in order to be used in a durable execution environment like Temporal. The ID will be used to identify the toolset's activities within the workflow."
+        )
+
+    settings = settings or TemporalSettings()
+
+    original_call_tool = toolset.call_tool
+
+    @activity.defn(name=f'function_toolset__{id}__call_tool')
+    async def call_tool_activity(params: _CallToolParams) -> Any:
+        name = params.name
+        ctx = settings.for_tool(id, name).deserialize_run_context(params.serialized_run_context)
+        tool = (await toolset.get_tools(ctx))[name]
+        return await original_call_tool(name, params.tool_args, ctx, tool)
+
+    async def call_tool(name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
+        tool_settings = settings.for_tool(id, name)
+        serialized_run_context = tool_settings.serialize_run_context(ctx)
+        return await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+            activity=call_tool_activity,
+            arg=_CallToolParams(name=name, tool_args=tool_args, serialized_run_context=serialized_run_context),
+            **tool_settings.execute_activity_kwargs,
+        )
+
+    toolset.call_tool = call_tool
+
+    activities = [call_tool_activity]
+    setattr(toolset, '__temporal_activities', activities)
+    return activities
+
+
+# class TemporalFunctionToolset(FunctionToolset[AgentDepsT]):
+#     def __init__(
+#         self,
+#         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
+#         max_retries: int = 1,
+#         temporal_settings: TemporalSettings | None = None,
+#         serialize_run_context: Callable[[RunContext[AgentDepsT]], Any] | None = None,
+#         deserialize_run_context: Callable[[Any], RunContext[AgentDepsT]] | None = None,
+#     ):
+#         super().__init__(tools, max_retries)
+#         self.temporal_settings = temporal_settings or TemporalSettings()
+#         self.serialize_run_context = serialize_run_context or TemporalRunContext[AgentDepsT].serialize_run_context
+#         self.deserialize_run_context = deserialize_run_context or TemporalRunContext[AgentDepsT].deserialize_run_context
+
+#         @activity.defn(name='function_toolset_call_tool')
+#         async def call_tool_activity(params: FunctionCallToolParams) -> Any:
+#             ctx = self.deserialize_run_context(params.serialized_run_context)
+#             tool = (await self.get_tools(ctx))[params.name]
+#             return await FunctionToolset[AgentDepsT].call_tool(self, params.name, params.tool_args, ctx, tool)
+
+#         self.call_tool_activity = call_tool_activity
+
+#     async def call_tool(
+#         self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+#     ) -> Any:
+#         serialized_run_context = self.serialize_run_context(ctx)
+#         return await workflow.execute_activity(
+#             activity=self.call_tool_activity,
+#             arg=FunctionCallToolParams(name=name, tool_args=tool_args, serialized_run_context=serialized_run_context),
+#             **self.temporal_settings.__dict__,
+#         )