refactor shared code into shared location

Author: wreed4

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py

@@ -1,57 +1,25 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Annotated, Any, Literal
+from typing import Any, Literal
 
-from pydantic import ConfigDict, Discriminator, with_config
 from temporalio import activity, workflow
 from temporalio.workflow import ActivityConfig
-from typing_extensions import assert_never
 
 from pydantic_ai import FunctionToolset, ToolsetTool
-from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.tools import AgentDepsT, RunContext
 from pydantic_ai.toolsets.function import FunctionToolsetTool
 
 from ._run_context import TemporalRunContext
-from ._toolset import TemporalWrapperToolset
-
-
-@dataclass
-@with_config(ConfigDict(arbitrary_types_allowed=True))
-class _CallToolParams:
-    name: str
-    tool_args: dict[str, Any]
-    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'),
-]
+from ._toolset import (
+    TemporalWrapperToolset,
+    _CallToolParams,
+    _CallToolResult,
+    _ToolReturn,
+    remap_dataclass_to_exception,
+    remap_exception_to_dataclass,
+)
 
 
 class TemporalFunctionToolset(TemporalWrapperToolset[AgentDepsT]):
@@ -87,12 +55,8 @@ async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _Call
             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)
+            except Exception as e:
+                return remap_exception_to_dataclass(e)
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         call_tool_activity.__annotations__['deps'] = deps_type
@@ -130,18 +94,10 @@ async def call_tool(
                     name=name,
                     tool_args=tool_args,
                     serialized_run_context=serialized_run_context,
+                    tool_def=None,
                 ),
                 ctx.deps,
             ],
             **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)
+        return remap_dataclass_to_exception(result)

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp_server.py

@@ -1,64 +1,27 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Annotated, Any, Literal
+from typing import Any, Literal
 
-from pydantic import ConfigDict, Discriminator, with_config
 from temporalio import activity, workflow
 from temporalio.workflow import ActivityConfig
 from typing_extensions import Self
 
 from pydantic_ai import ToolsetTool
-from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
+from pydantic_ai.exceptions import UserError
 from pydantic_ai.mcp import MCPServer, ToolResult
 from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
 
 from ._run_context import TemporalRunContext
-from ._toolset import TemporalWrapperToolset
-
-
-@dataclass
-@with_config(ConfigDict(arbitrary_types_allowed=True))
-class _GetToolsParams:
-    serialized_run_context: Any
-
-
-@dataclass
-@with_config(ConfigDict(arbitrary_types_allowed=True))
-class _CallToolParams:
-    name: str
-    tool_args: dict[str, Any]
-    serialized_run_context: Any
-    tool_def: ToolDefinition
-
-
-@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'),
-]
+from ._toolset import (
+    TemporalWrapperToolset,
+    _CallToolParams,
+    _CallToolResult,
+    _GetToolsParams,
+    _ToolReturn,
+    remap_exception_to_dataclass,
+    remap_dataclass_to_exception,
+)
 
 
 class TemporalMCPServer(TemporalWrapperToolset[AgentDepsT]):
@@ -103,19 +66,16 @@ async def get_tools_activity(params: _GetToolsParams, deps: AgentDepsT) -> dict[
         async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _CallToolResult:
             run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             try:
+                assert isinstance(params.tool_def, ToolDefinition)
                 result = await self.wrapped.call_tool(
                     params.name,
                     params.tool_args,
                     run_context,
                     self.tool_for_tool_def(params.tool_def),
                 )
                 return _ToolReturn(result=result)
-            except ApprovalRequired:
-                return _ApprovalRequired()
-            except CallDeferred:
-                return _CallDeferred()
-            except ModelRetry as e:
-                return _ModelRetry(message=e.message)
+            except Exception as e:
+                return remap_exception_to_dataclass(e)
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         call_tool_activity.__annotations__['deps'] = deps_type
@@ -161,13 +121,13 @@ async def call_tool(
         tool_args: dict[str, Any],
         ctx: RunContext[AgentDepsT],
         tool: ToolsetTool[AgentDepsT],
-    ) -> ToolResult:
+    ) -> _CallToolResult:
         if not workflow.in_workflow():
             return await super().call_tool(name, tool_args, ctx, tool)
 
         tool_activity_config = self.activity_config | self.tool_activity_config.get(name, {})
         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(
@@ -180,3 +140,4 @@ async def call_tool(
             ],
             **tool_activity_config,
         )
+        return remap_dataclass_to_exception(result)

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_toolset.py

@@ -2,16 +2,63 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import Any, Literal
+from dataclasses import dataclass
+from typing import Annotated, Any, Literal
 
+from pydantic import ConfigDict, Discriminator, with_config
 from temporalio.workflow import ActivityConfig
+from typing_extensions import assert_never
 
 from pydantic_ai import AbstractToolset, FunctionToolset, WrapperToolset
-from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
+from pydantic_ai.tools import AgentDepsT, ToolDefinition
 
 from ._run_context import TemporalRunContext
 
 
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class _GetToolsParams:
+    serialized_run_context: Any
+
+
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class _CallToolParams:
+    name: str
+    tool_args: dict[str, Any]
+    serialized_run_context: Any
+    tool_def: ToolDefinition | None
+
+
+@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 TemporalWrapperToolset(WrapperToolset[AgentDepsT], ABC):
     @property
     def id(self) -> str:
@@ -31,6 +78,30 @@ def visit_and_replace(
         return self
 
 
+def remap_exception_to_dataclass(e: Exception) -> _CallToolResult:
+    try:
+        raise e
+    except ApprovalRequired:
+        return _ApprovalRequired()
+    except CallDeferred:
+        return _CallDeferred()
+    except ModelRetry as e:
+        return _ModelRetry(message=e.message)
+
+
+def remap_dataclass_to_exception(o: _CallToolResult):
+    if isinstance(o, _ApprovalRequired):
+        raise ApprovalRequired()
+    elif isinstance(o, _CallDeferred):
+        raise CallDeferred()
+    elif isinstance(o, _ModelRetry):
+        raise ModelRetry(o.message)
+    elif isinstance(o, _ToolReturn):
+        return o.result
+    else:
+        assert_never(o)
+
+
 def temporalize_toolset(
     toolset: AbstractToolset[AgentDepsT],
     activity_name_prefix: str,