Update agent.py

Author: DanielHashmi

src/agents/agent.py

@@ -10,7 +10,7 @@
 import asyncio
 import inspect
 from collections.abc import Awaitable
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, cast
 
 from openai.types.responses.response_prompt_param import ResponsePromptParam
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
@@ -36,29 +36,25 @@
 class ToolsToFinalOutputResult(BaseModel):
     """
     Result object for tool-to-final output processing.
-    
+
     Attributes:
         is_final_output: Whether this is the final output. If False, the LLM will run again.
         final_output: The final output value, if is_final_output is True.
     """
-    
+
     model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-        validate_assignment=True,
-        extra='forbid',
-        frozen=True
+        arbitrary_types_allowed=True, validate_assignment=True, extra="forbid", frozen=True
     )
-    
+
     is_final_output: bool = Field(
         ...,  # Required field
-        description="Whether this is the final output. If False, the LLM will run again."
+        description="Whether this is the final output. If False, the LLM will run again.",
     )
     final_output: Any | None = Field(
-        default=None,
-        description="The final output value, if is_final_output is True."
+        default=None, description="The final output value, if is_final_output is True."
     )
-    
-    @field_validator('is_final_output')
+
+    @field_validator("is_final_output")
     @classmethod
     def validate_is_final_output(cls, v: Any) -> bool:
         """Validate is_final_output is a boolean value."""
@@ -79,27 +75,29 @@ def validate_is_final_output(cls, v: Any) -> bool:
 
 class StopAtTools(TypedDict):
     """Configuration for stopping agent execution at specific tools."""
+
     stop_at_tool_names: list[str]
 
 
 class MCPConfig(TypedDict):
     """Configuration for Model Context Protocol servers."""
+
     convert_schemas_to_strict: NotRequired[bool]
 
 
 def is_stop_at_tools_dict(v: dict) -> bool:
     """
     Validate if a dictionary matches the StopAtTools structure.
-    
+
     Args:
         v: Dictionary to validate
-        
+
     Returns:
         bool: True if dictionary is valid StopAtTools structure
     """
     return (
-        isinstance(v, dict) 
-        and "stop_at_tool_names" in v 
+        isinstance(v, dict)
+        and "stop_at_tool_names" in v
         and isinstance(v["stop_at_tool_names"], list)
         and all(isinstance(x, str) for x in v["stop_at_tool_names"])
     )
@@ -108,77 +106,71 @@ def is_stop_at_tools_dict(v: dict) -> bool:
 class AgentBase(BaseModel, Generic[TContext]):
     """
     Base class for Agent implementations providing core functionality.
-    
+
     This class implements the base agent functionality including tool management,
     MCP server configuration, and validation.
-    
+
     Attributes:
         name: The name of the agent
         handoff_description: Optional description for handoff functionality
         tools: List of available tools
         mcp_servers: List of MCP servers
         mcp_config: Configuration for MCP servers
     """
-    
+
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
         validate_assignment=True,
-        extra='forbid',
+        extra="forbid",
         frozen=True,
-        defer_build=True
+        defer_build=True,
     )
 
     name: str = Field(
         ...,  # Required field
-        description="The name of the agent."
+        description="The name of the agent.",
     )
     handoff_description: str | None = Field(
-        default=None,
-        description="Description used when the agent is used as a handoff."
+        default=None, description="Description used when the agent is used as a handoff."
     )
     tools: list[Tool] = Field(
-        default_factory=list,
-        description="List of tools available to the agent."
+        default_factory=list, description="List of tools available to the agent."
     )
     mcp_servers: list[Any] = Field(
-        default_factory=list,
-        description="List of MCP servers available to the agent."
+        default_factory=list, description="List of MCP servers available to the agent."
     )
     mcp_config: dict[str, Any] = Field(
         default_factory=lambda: {"convert_schemas_to_strict": False},
-        description="MCP configuration settings."
+        description="MCP configuration settings.",
     )
 
     # ... (validators remain the same but with improved docstrings) ...
 
     async def get_mcp_tools(self, run_context: RunContextWrapper[TContext]) -> list[Tool]:
         """
         Fetch available tools from MCP servers.
-        
+
         Args:
             run_context: Current run context wrapper
-            
+
         Returns:
             list[Tool]: List of available MCP tools
         """
         convert_schemas_to_strict = self.mcp_config.get("convert_schemas_to_strict", False)
         if not self.mcp_servers:
             return []
-        
+
         return await MCPUtil.get_all_function_tools(
-            self.mcp_servers,
-            convert_schemas_to_strict,
-            run_context,
-            self
+            self.mcp_servers, convert_schemas_to_strict, run_context, self
         )
 
     async def get_all_tools(self, run_context: RunContextWrapper[Any]) -> list[Tool]:
         """
         Get all available tools, including MCP tools and function tools.
-        
+
         Args:
             run_context: Current run context wrapper
-            
+
         Returns:
             list[Tool]: Combined list of available tools
         """
@@ -205,10 +197,10 @@ async def _check_tool_enabled(tool: Tool) -> bool:
 class Agent(AgentBase, Generic[TContext]):
     """
     Primary agent implementation with full configuration and behavior support.
-    
+
     This class extends AgentBase with additional functionality for instructions,
     prompts, guardrails, and tool behavior configuration.
-    
+
     Attributes:
         instructions: Agent instructions/system prompt
         prompt: Prompt configuration
@@ -222,94 +214,73 @@ class Agent(AgentBase, Generic[TContext]):
         tool_use_behavior: Tool usage configuration
         reset_tool_choice: Tool choice reset behavior
     """
-    
+
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
         validate_assignment=True,
-        extra='forbid',
+        extra="forbid",
         frozen=True,
-        defer_build=True
+        defer_build=True,
     )
 
     instructions: (
         str
         | Callable[
-            [RunContextWrapper[TContext], 'Agent[TContext]'],
+            [RunContextWrapper[TContext], Agent[TContext]],
             MaybeAwaitable[str],
         ]
         | None
-    ) = Field(
-        default=None,
-        description="The instructions/system prompt for the agent."
-    )
+    ) = Field(default=None, description="The instructions/system prompt for the agent.")
 
     prompt: Prompt | DynamicPromptFunction | None = Field(
-        default=None,
-        description="Prompt configuration for the agent."
+        default=None, description="Prompt configuration for the agent."
     )
 
     handoffs: list[Any] = Field(
-        default_factory=list,
-        description="List of sub-agents for delegation."
+        default_factory=list, description="List of sub-agents for delegation."
     )
 
-    model: str | Model | None = Field(
-        default=None,
-        description="Model implementation to use."
-    )
+    model: str | Model | None = Field(default=None, description="Model implementation to use.")
 
     model_settings: ModelSettings = Field(
-        default_factory=ModelSettings,
-        description="Model-specific configuration."
+        default_factory=ModelSettings, description="Model-specific configuration."
     )
 
     input_guardrails: list[InputGuardrail[TContext]] = Field(
-        default_factory=list,
-        description="Pre-execution validation checks."
+        default_factory=list, description="Pre-execution validation checks."
     )
 
     output_guardrails: list[OutputGuardrail[TContext]] = Field(
-        default_factory=list,
-        description="Post-execution validation checks."
+        default_factory=list, description="Post-execution validation checks."
     )
 
     output_type: type[Any] | AgentOutputSchemaBase | None = Field(
-        default=None,
-        description="Output type specification."
+        default=None, description="Output type specification."
     )
 
-    hooks: Any | None = Field(
-        default=None,
-        description="Lifecycle event callbacks."
-    )
+    hooks: Any | None = Field(default=None, description="Lifecycle event callbacks.")
 
-    tool_use_behavior: Union[
-        Literal["run_llm_again", "stop_on_first_tool"],
-        StopAtTools,
-        ToolsToFinalOutputFunction
-    ] = Field(
-        default="run_llm_again",
-        description="Tool usage behavior configuration."
-    )
+    tool_use_behavior: (
+        Literal["run_llm_again", "stop_on_first_tool"] | StopAtTools | ToolsToFinalOutputFunction
+    ) = Field(default="run_llm_again", description="Tool usage behavior configuration.")
 
     reset_tool_choice: bool = Field(
-        default=True,
-        description="Whether to reset tool choice after use."
+        default=True, description="Whether to reset tool choice after use."
     )
 
     # ... (validators remain the same but with improved docstrings) ...
 
-    @model_validator(mode='after')
-    def validate_model_configuration(self) -> 'Agent':
+    @model_validator(mode="after")
+    def validate_model_configuration(self) -> Agent:
         """
         Validate complete model configuration.
-        
+
         This validator ensures that the model configuration is consistent,
         particularly the relationship between model and prompt settings.
-        
+
         Returns:
             Agent: The validated agent instance
-            
+
         Raises:
             ValueError: If validation fails
         """
@@ -320,13 +291,13 @@ def validate_model_configuration(self) -> 'Agent':
     def clone(self, **kwargs: Any) -> Agent[TContext]:
         """
         Create a copy of the agent with specified modifications.
-        
+
         Args:
             **kwargs: Fields to override in the new instance
-            
+
         Returns:
             Agent[TContext]: New agent instance with specified changes
-            
+
         Note:
             This performs a shallow copy. Mutable attributes like tools and
             handoffs are shallow-copied.
@@ -341,15 +312,16 @@ def as_tool(
     ) -> Tool:
         """
         Convert this agent into a tool callable by other agents.
-        
+
         Args:
             tool_name: Name for the tool (defaults to agent name)
             tool_description: Description of the tool's functionality
             custom_output_extractor: Custom function to extract output
-            
+
         Returns:
             Tool: Tool instance wrapping this agent
         """
+
         @function_tool(
             name_override=tool_name or _transforms.transform_string_function_style(self.name),
             description_override=tool_description or "",
@@ -372,13 +344,13 @@ async def run_agent(context: RunContextWrapper, input: str) -> str:
     async def get_system_prompt(self, run_context: RunContextWrapper[TContext]) -> str | None:
         """
         Get the system prompt for the agent.
-        
+
         Args:
             run_context: Current run context
-            
+
         Returns:
             str | None: System prompt if available
-            
+
         Note:
             Handles both static strings and dynamic prompt generation.
         """
@@ -397,10 +369,10 @@ async def get_prompt(
     ) -> ResponsePromptParam | None:
         """
         Get the prompt configuration for the agent.
-        
+
         Args:
             run_context: Current run context
-            
+
         Returns:
             ResponsePromptParam | None: Prompt configuration if available
         """