Use pydantic serialization for ActivityModelInput

Author: AnatoliB

samples-v2/openai_agents/model_invoker.py

@@ -4,6 +4,7 @@
 from datetime import timedelta
 from typing import Any, Optional, Union
 
+from pydantic import BaseModel, Field
 from agents import (
     AgentOutputSchemaBase,
     CodeInterpreterTool,
@@ -28,7 +29,6 @@
     AsyncOpenAI,
 )
 from openai.types.responses.tool_param import Mcp
-from pydantic_core import to_json
 try:
     from azure.durable_functions import ApplicationError
 except ImportError:
@@ -40,8 +40,7 @@ def __init__(self, message: str, non_retryable: bool = False, next_retry_delay =
             self.next_retry_delay = next_retry_delay
 
 
-@dataclass
-class HandoffInput:
+class HandoffInput(BaseModel):
     """Data conversion friendly representation of a Handoff. Contains only the fields which are needed by the model
     execution to determine what to handoff to, not the actual handoff invocation, which remains in the workflow context.
     """
@@ -53,8 +52,7 @@ class HandoffInput:
     strict_json_schema: bool = True
 
 
-@dataclass
-class FunctionToolInput:
+class FunctionToolInput(BaseModel):
     """Data conversion friendly representation of a FunctionTool. Contains only the fields which are needed by the model
     execution to determine what tool to call, not the actual tool invocation, which remains in the workflow context.
     """
@@ -65,8 +63,7 @@ class FunctionToolInput:
     strict_json_schema: bool = True
 
 
-@dataclass
-class HostedMCPToolInput:
+class HostedMCPToolInput(BaseModel):
     """Data conversion friendly representation of a HostedMCPTool. Contains only the fields which are needed by the model
     execution to determine what tool to call, not the actual tool invocation, which remains in the workflow context.
     """
@@ -84,13 +81,12 @@ class HostedMCPToolInput:
 ]
 
 
-@dataclass
-class AgentOutputSchemaInput(AgentOutputSchemaBase):
+class AgentOutputSchemaInput(AgentOutputSchemaBase, BaseModel):
     """Data conversion friendly representation of AgentOutputSchema."""
 
-    output_type_name: Optional[str]
+    output_type_name: Optional[str] = None
     is_wrapped: bool
-    output_schema: Optional[dict[str, Any]]
+    output_schema: Optional[dict[str, Any]] = None
     strict_json_schema: bool
 
     def is_plain_text(self) -> bool:
@@ -131,74 +127,28 @@ class ModelTracingInput(enum.IntEnum):
     ENABLED_WITHOUT_DATA = 2
 
 
-@dataclass
-class ActivityModelInput:
+class ActivityModelInput(BaseModel):
     """Input for the invoke_model_activity activity."""
 
     input: Union[str, list[TResponseInputItem]]
     model_settings: ModelSettings
     tracing: ModelTracingInput
     model_name: Optional[str] = None
     system_instructions: Optional[str] = None
-    tools: list[ToolInput] = None
+    tools: list[ToolInput] = Field(default_factory=list)
     output_schema: Optional[AgentOutputSchemaInput] = None
-    handoffs: list[HandoffInput] = None
+    handoffs: list[HandoffInput] = Field(default_factory=list)
     previous_response_id: Optional[str] = None
     prompt: Optional[Any] = None
 
-    def __post_init__(self):
-        """Initialize default values for list fields."""
-        if self.tools is None:
-            self.tools = []
-        if self.handoffs is None:
-            self.handoffs = []
-
     def to_json(self) -> str:
         """Convert the ActivityModelInput to a JSON string."""
-        return to_json(self).decode('utf-8')
+        return self.model_dump_json()
 
     @classmethod
     def from_json(cls, json_str: str) -> 'ActivityModelInput':
         """Create an ActivityModelInput instance from a JSON string."""
-        import json
-        data = json.loads(json_str)
-        
-        # Convert complex types back from dictionaries
-        if 'model_settings' in data and isinstance(data['model_settings'], dict):
-            data['model_settings'] = ModelSettings(**data['model_settings'])
-        
-        if 'tracing' in data and isinstance(data['tracing'], int):
-            data['tracing'] = ModelTracingInput(data['tracing'])
-        
-        # Convert tool inputs back to proper types
-        if 'tools' in data and data['tools']:
-            converted_tools = []
-            for tool_data in data['tools']:
-                if isinstance(tool_data, dict):
-                    # Check the tool type and convert accordingly
-                    if 'name' in tool_data and 'description' in tool_data and 'params_json_schema' in tool_data:
-                        # FunctionToolInput
-                        converted_tools.append(FunctionToolInput(**tool_data))
-                    elif 'tool_config' in tool_data:
-                        # HostedMCPToolInput
-                        converted_tools.append(HostedMCPToolInput(**tool_data))
-                    else:
-                        # For other tool types like FileSearchTool, WebSearchTool, etc.
-                        # These might be already properly serialized/deserialized by pydantic_core
-                        converted_tools.append(tool_data)
-                else:
-                    converted_tools.append(tool_data)
-            data['tools'] = converted_tools
-        
-        # Convert handoffs back to proper types
-        if 'handoffs' in data and data['handoffs']:
-            data['handoffs'] = [HandoffInput(**handoff) for handoff in data['handoffs']]
-        
-        # Convert output_schema back to proper type
-        if 'output_schema' in data and data['output_schema'] and isinstance(data['output_schema'], dict):
-            data['output_schema'] = AgentOutputSchemaInput(**data['output_schema'])
-        
-        return cls(**data)
+        return cls.model_validate_json(json_str)
 
 
 class ModelInvoker:
@@ -220,7 +170,7 @@ async def empty_on_invoke_handoff(
 
         # workaround for https://github.com/pydantic/pydantic/issues/9541
         # ValidatorIterator returned
-        input_json = to_json(input.input)
+        input_json = json.dumps(input.input, default=str)
         input_input = json.loads(input_json)
 
         def make_tool(tool: ToolInput) -> Tool: