rename to make classes non-private

Author: wreed4

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_function_toolset.py

@@ -14,9 +14,9 @@
 from ._run_context import TemporalRunContext
 from ._toolset import (
     TemporalWrapperToolset,
-    _CallToolParams,
-    _CallToolResult,
-    _ToolReturn,
+    CallToolParamsData,
+    CallToolResultData,
+    ToolReturnData,
     remap_dataclass_to_exception,
     remap_exception_to_dataclass,
 )
@@ -38,7 +38,7 @@ def __init__(
         self.tool_activity_config = tool_activity_config
         self.run_context_type = run_context_type
 
-        async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _CallToolResult:
+        async def call_tool_activity(params: CallToolParamsData, deps: AgentDepsT) -> CallToolResultData:
             name = params.name
             ctx = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             try:
@@ -54,7 +54,7 @@ async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _Call
             args_dict = tool.args_validator.validate_python(params.tool_args)
             try:
                 result = await self.wrapped.call_tool(name, args_dict, ctx, tool)
-                return _ToolReturn(result=result)
+                return ToolReturnData(result=result)
             except Exception as e:
                 return remap_exception_to_dataclass(e)
 
@@ -90,7 +90,7 @@ async def call_tool(
         result = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.call_tool_activity,
             args=[
-                _CallToolParams(
+                CallToolParamsData(
                     name=name,
                     tool_args=tool_args,
                     serialized_run_context=serialized_run_context,

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp_server.py

@@ -15,10 +15,10 @@
 from ._run_context import TemporalRunContext
 from ._toolset import (
     TemporalWrapperToolset,
-    _CallToolParams,
-    _CallToolResult,
-    _GetToolsParams,
-    _ToolReturn,
+    CallToolParamsData,
+    CallToolResultData,
+    GetToolsParamsData,
+    ToolReturnData,
     remap_exception_to_dataclass,
     remap_dataclass_to_exception,
 )
@@ -49,7 +49,7 @@ def __init__(
 
         self.run_context_type = run_context_type
 
-        async def get_tools_activity(params: _GetToolsParams, deps: AgentDepsT) -> dict[str, ToolDefinition]:
+        async def get_tools_activity(params: GetToolsParamsData, 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,7 +63,7 @@ async def get_tools_activity(params: _GetToolsParams, deps: AgentDepsT) -> dict[
             get_tools_activity
         )
 
-        async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _CallToolResult:
+        async def call_tool_activity(params: CallToolParamsData, deps: AgentDepsT) -> CallToolResultData:
             run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             try:
                 assert isinstance(params.tool_def, ToolDefinition)
@@ -73,7 +73,7 @@ async def call_tool_activity(params: _CallToolParams, deps: AgentDepsT) -> _Call
                     run_context,
                     self.tool_for_tool_def(params.tool_def),
                 )
-                return _ToolReturn(result=result)
+                return ToolReturnData(result=result)
             except Exception as e:
                 return remap_exception_to_dataclass(e)
 
@@ -108,7 +108,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=[
-                _GetToolsParams(serialized_run_context=serialized_run_context),
+                GetToolsParamsData(serialized_run_context=serialized_run_context),
                 ctx.deps,
             ],
             **self.activity_config,
@@ -121,7 +121,7 @@ async def call_tool(
         tool_args: dict[str, Any],
         ctx: RunContext[AgentDepsT],
         tool: ToolsetTool[AgentDepsT],
-    ) -> _CallToolResult:
+    ) -> CallToolResultData:
         if not workflow.in_workflow():
             return await super().call_tool(name, tool_args, ctx, tool)
 
@@ -130,7 +130,7 @@ async def call_tool(
         result = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
             activity=self.call_tool_activity,
             args=[
-                _CallToolParams(
+                CallToolParamsData(
                     name=name,
                     tool_args=tool_args,
                     serialized_run_context=serialized_run_context,

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_toolset.py

@@ -18,43 +18,43 @@
 
 @dataclass
 @with_config(ConfigDict(arbitrary_types_allowed=True))
-class _GetToolsParams:
+class GetToolsParamsData:
     serialized_run_context: Any
 
 
 @dataclass
 @with_config(ConfigDict(arbitrary_types_allowed=True))
-class _CallToolParams:
+class CallToolParamsData:
     name: str
     tool_args: dict[str, Any]
     serialized_run_context: Any
     tool_def: ToolDefinition | None
 
 
 @dataclass
-class _ApprovalRequired:
+class ApprovalRequiredData:
     kind: Literal['approval_required'] = 'approval_required'
 
 
 @dataclass
-class _CallDeferred:
+class CallDeferredData:
     kind: Literal['call_deferred'] = 'call_deferred'
 
 
 @dataclass
-class _ModelRetry:
+class ModelRetryData:
     message: str
     kind: Literal['model_retry'] = 'model_retry'
 
 
 @dataclass
-class _ToolReturn:
+class ToolReturnData:
     result: Any
     kind: Literal['tool_return'] = 'tool_return'
 
 
-_CallToolResult = Annotated[
-    _ApprovalRequired | _CallDeferred | _ModelRetry | _ToolReturn,
+CallToolResultData = Annotated[
+    ApprovalRequiredData | CallDeferredData | ModelRetryData | ToolReturnData,
     Discriminator('kind'),
 ]
 
@@ -78,25 +78,25 @@ def visit_and_replace(
         return self
 
 
-def remap_exception_to_dataclass(e: Exception) -> _CallToolResult:
+def remap_exception_to_dataclass(e: Exception) -> CallToolResultData:
     try:
         raise e
     except ApprovalRequired:
-        return _ApprovalRequired()
+        return ApprovalRequiredData()
     except CallDeferred:
-        return _CallDeferred()
+        return CallDeferredData()
     except ModelRetry as e:
-        return _ModelRetry(message=e.message)
+        return ModelRetryData(message=e.message)
 
 
-def remap_dataclass_to_exception(o: _CallToolResult):
-    if isinstance(o, _ApprovalRequired):
+def remap_dataclass_to_exception(o: CallToolResultData):
+    if isinstance(o, ApprovalRequiredData):
         raise ApprovalRequired()
-    elif isinstance(o, _CallDeferred):
+    elif isinstance(o, CallDeferredData):
         raise CallDeferred()
-    elif isinstance(o, _ModelRetry):
+    elif isinstance(o, ModelRetryData):
         raise ModelRetry(o.message)
-    elif isinstance(o, _ToolReturn):
+    elif isinstance(o, ToolReturnData):
         return o.result
     else:
         assert_never(o)