move some stuff around

Author: DouweM

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py

@@ -13,12 +13,9 @@
 
 from ._run_context import TemporalRunContext
 from ._toolset import (
-    CallToolParamsData,
-    CallToolResultData,
+    CallToolParams,
+    CallToolResult,
     TemporalWrapperToolset,
-    ToolReturnData,
-    remap_dataclass_to_exception,
-    remap_exception_to_dataclass,
 )
 
 
@@ -38,7 +35,7 @@ def __init__(
         self.tool_activity_config = tool_activity_config
         self.run_context_type = run_context_type
 
-        async def call_tool_activity(params: CallToolParamsData, deps: AgentDepsT) -> CallToolResultData:
+        async def call_tool_activity(params: CallToolParams, deps: AgentDepsT) -> CallToolResult:
             name = params.name
             ctx = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             try:
@@ -52,11 +49,7 @@ async def call_tool_activity(params: CallToolParamsData, deps: AgentDepsT) -> Ca
             # The tool args will already have been validated into their proper types in the `ToolManager`,
             # but `execute_activity` would have turned them into simple Python types again, so we need to re-validate them.
             args_dict = tool.args_validator.validate_python(params.tool_args)
-            try:
-                result = await self.wrapped.call_tool(name, args_dict, ctx, tool)
-                return ToolReturnData(result=result)
-            except Exception as e:
-                return remap_exception_to_dataclass(e)
+            return await self._wrap_call_tool_result(self.wrapped.call_tool(name, args_dict, ctx, tool))
 
         # Set type hint explicitly so that Temporal can take care of serialization and deserialization
         call_tool_activity.__annotations__['deps'] = deps_type
@@ -87,17 +80,18 @@ async def call_tool(
 
         tool_activity_config = self.activity_config | tool_activity_config
         serialized_run_context = self.run_context_type.serialize_run_context(ctx)
-        result = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
-            activity=self.call_tool_activity,
-            args=[
-                CallToolParamsData(
-                    name=name,
-                    tool_args=tool_args,
-                    serialized_run_context=serialized_run_context,
-                    tool_def=None,
-                ),
-                ctx.deps,
-            ],
-            **tool_activity_config,
+        return self._unwrap_call_tool_result(
+            await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+                activity=self.call_tool_activity,
+                args=[
+                    CallToolParams(
+                        name=name,
+                        tool_args=tool_args,
+                        serialized_run_context=serialized_run_context,
+                        tool_def=None,
+                    ),
+                    ctx.deps,
+                ],
+                **tool_activity_config,
+            )
         )
-        return remap_dataclass_to_exception(result)

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp_server.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import Any, Literal
 
+from pydantic import ConfigDict, with_config
 from temporalio import activity, workflow
 from temporalio.workflow import ActivityConfig
 from typing_extensions import Self
@@ -14,16 +16,18 @@
 
 from ._run_context import TemporalRunContext
 from ._toolset import (
-    CallToolParamsData,
-    CallToolResultData,
-    GetToolsParamsData,
+    CallToolParams,
+    CallToolResult,
     TemporalWrapperToolset,
-    ToolReturnData,
-    remap_dataclass_to_exception,
-    remap_exception_to_dataclass,
 )
 
 
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class _GetToolsParams:
+    serialized_run_context: Any
+
+
 class TemporalMCPServer(TemporalWrapperToolset[AgentDepsT]):
     def __init__(
         self,
@@ -49,7 +53,7 @@ def __init__(
 
         self.run_context_type = run_context_type
 
-        async def get_tools_activity(params: GetToolsParamsData, deps: AgentDepsT) -> dict[str, ToolDefinition]:
+        async def get_tools_activity(params: _GetToolsParams, deps: AgentDepsT) -> dict[str, ToolDefinition]:
             run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             tools = await self.wrapped.get_tools(run_context)
             # ToolsetTool is not serializable as it holds a SchemaValidator (which is also the same for every MCP tool so unnecessary to pass along the wire every time),
@@ -63,19 +67,17 @@ async def get_tools_activity(params: GetToolsParamsData, deps: AgentDepsT) -> di
             get_tools_activity
         )
 
-        async def call_tool_activity(params: CallToolParamsData, deps: AgentDepsT) -> CallToolResultData:
+        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(
+            assert isinstance(params.tool_def, ToolDefinition)
+            return await self._wrap_call_tool_result(
+                self.wrapped.call_tool(
                     params.name,
                     params.tool_args,
                     run_context,
                     self.tool_for_tool_def(params.tool_def),
                 )
-                return ToolReturnData(result=result)
-            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
@@ -108,7 +110,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[
         tool_defs = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.get_tools_activity,
             args=[
-                GetToolsParamsData(serialized_run_context=serialized_run_context),
+                _GetToolsParams(serialized_run_context=serialized_run_context),
                 ctx.deps,
             ],
             **self.activity_config,
@@ -121,23 +123,24 @@ async def call_tool(
         tool_args: dict[str, Any],
         ctx: RunContext[AgentDepsT],
         tool: ToolsetTool[AgentDepsT],
-    ) -> CallToolResultData:
+    ) -> 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)
-        result = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
-            activity=self.call_tool_activity,
-            args=[
-                CallToolParamsData(
-                    name=name,
-                    tool_args=tool_args,
-                    serialized_run_context=serialized_run_context,
-                    tool_def=tool.tool_def,
-                ),
-                ctx.deps,
-            ],
-            **tool_activity_config,
+        return self._unwrap_call_tool_result(
+            await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+                activity=self.call_tool_activity,
+                args=[
+                    CallToolParams(
+                        name=name,
+                        tool_args=tool_args,
+                        serialized_run_context=serialized_run_context,
+                        tool_def=tool.tool_def,
+                    ),
+                    ctx.deps,
+                ],
+                **tool_activity_config,
+            )
         )
-        return remap_dataclass_to_exception(result)

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_toolset.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import Annotated, Any, Literal
 
@@ -18,43 +18,37 @@
 
 @dataclass
 @with_config(ConfigDict(arbitrary_types_allowed=True))
-class GetToolsParamsData:
-    serialized_run_context: Any
-
-
-@dataclass
-@with_config(ConfigDict(arbitrary_types_allowed=True))
-class CallToolParamsData:
+class CallToolParams:
     name: str
     tool_args: dict[str, Any]
     serialized_run_context: Any
     tool_def: ToolDefinition | None
 
 
 @dataclass
-class ApprovalRequiredData:
+class _ApprovalRequired:
     kind: Literal['approval_required'] = 'approval_required'
 
 
 @dataclass
-class CallDeferredData:
+class _CallDeferred:
     kind: Literal['call_deferred'] = 'call_deferred'
 
 
 @dataclass
-class ModelRetryData:
+class _ModelRetry:
     message: str
     kind: Literal['model_retry'] = 'model_retry'
 
 
 @dataclass
-class ToolReturnData:
+class _ToolReturn:
     result: Any
     kind: Literal['tool_return'] = 'tool_return'
 
 
-CallToolResultData = Annotated[
-    ApprovalRequiredData | CallDeferredData | ModelRetryData | ToolReturnData,
+CallToolResult = Annotated[
+    _ApprovalRequired | _CallDeferred | _ModelRetry | _ToolReturn,
     Discriminator('kind'),
 ]
 
@@ -77,29 +71,28 @@ def visit_and_replace(
         # Temporalized toolsets cannot be swapped out after the fact.
         return self
 
-
-def remap_exception_to_dataclass(e: Exception) -> CallToolResultData:
-    try:
-        raise e
-    except ApprovalRequired:
-        return ApprovalRequiredData()
-    except CallDeferred:
-        return CallDeferredData()
-    except ModelRetry as e:
-        return ModelRetryData(message=e.message)
-
-
-def remap_dataclass_to_exception(o: CallToolResultData):
-    if isinstance(o, ApprovalRequiredData):
-        raise ApprovalRequired()
-    elif isinstance(o, CallDeferredData):
-        raise CallDeferred()
-    elif isinstance(o, ModelRetryData):
-        raise ModelRetry(o.message)
-    elif isinstance(o, ToolReturnData):
-        return o.result
-    else:
-        assert_never(o)
+    async def _wrap_call_tool_result(self, coro: Awaitable[Any]) -> CallToolResult:
+        try:
+            result = await coro
+            return _ToolReturn(result=result)
+        except ApprovalRequired:
+            return _ApprovalRequired()
+        except CallDeferred:
+            return _CallDeferred()
+        except ModelRetry as e:
+            return _ModelRetry(message=e.message)
+
+    def _unwrap_call_tool_result(self, result: CallToolResult) -> Any:
+        if isinstance(result, _ToolReturn):
+            return result.result
+        elif isinstance(result, _ApprovalRequired):
+            raise ApprovalRequired()
+        elif isinstance(result, _CallDeferred):
+            raise CallDeferred()
+        elif isinstance(result, _ModelRetry):
+            raise ModelRetry(result.message)
+        else:
+            assert_never(result)
 
 
 def temporalize_toolset(