openai-agents 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

agents/__init__.py

@@ -70,6 +70,7 @@ from .tracing import (
     GenerationSpanData,
     GuardrailSpanData,
     HandoffSpanData,
+    MCPListToolsSpanData,
     Span,
     SpanData,
     SpanError,
@@ -89,6 +90,7 @@ from .tracing import (
     get_current_trace,
     guardrail_span,
     handoff_span,
+    mcp_tools_span,
     set_trace_processors,
     set_tracing_disabled,
     set_tracing_export_api_key,
@@ -98,6 +100,7 @@ from .tracing import (
     transcription_span,
 )
 from .usage import Usage
+from .version import __version__
 
 
 def set_default_openai_key(key: str, use_for_tracing: bool = True) -> None:
@@ -220,6 +223,7 @@ __all__ = [
     "speech_group_span",
     "transcription_span",
     "speech_span",
+    "mcp_tools_span",
     "trace",
     "Trace",
     "TracingProcessor",
@@ -234,6 +238,7 @@ __all__ = [
     "HandoffSpanData",
     "SpeechGroupSpanData",
     "SpeechSpanData",
+    "MCPListToolsSpanData",
     "TranscriptionSpanData",
     "set_default_openai_key",
     "set_default_openai_client",
@@ -243,4 +248,5 @@ __all__ = [
     "gen_trace_id",
     "gen_span_id",
     "default_tool_error_function",
+    "__version__",
 ]

agents/_run_impl.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
 import inspect
 from collections.abc import Awaitable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, cast
 
 from openai.types.responses import (
@@ -47,10 +48,11 @@ from .items import (
 )
 from .lifecycle import RunHooks
 from .logger import logger
+from .model_settings import ModelSettings
 from .models.interface import ModelTracing
 from .run_context import RunContextWrapper, TContext
 from .stream_events import RunItemStreamEvent, StreamEvent
-from .tool import ComputerTool, FunctionTool, FunctionToolResult
+from .tool import ComputerTool, FunctionTool, FunctionToolResult, Tool
 from .tracing import (
     SpanError,
     Trace,
@@ -75,6 +77,23 @@ QUEUE_COMPLETE_SENTINEL = QueueCompleteSentinel()
 _NOT_FINAL_OUTPUT = ToolsToFinalOutputResult(is_final_output=False, final_output=None)
 
 
+@dataclass
+class AgentToolUseTracker:
+    agent_to_tools: list[tuple[Agent, list[str]]] = field(default_factory=list)
+    """Tuple of (agent, list of tools used). Can't use a dict because agents aren't hashable."""
+
+    def add_tool_use(self, agent: Agent[Any], tool_names: list[str]) -> None:
+        existing_data = next((item for item in self.agent_to_tools if item[0] == agent), None)
+        if existing_data:
+            existing_data[1].extend(tool_names)
+        else:
+            self.agent_to_tools.append((agent, tool_names))
+
+    def has_used_tools(self, agent: Agent[Any]) -> bool:
+        existing_data = next((item for item in self.agent_to_tools if item[0] == agent), None)
+        return existing_data is not None and len(existing_data[1]) > 0
+
+
 @dataclass
 class ToolRunHandoff:
     handoff: Handoff
@@ -99,6 +118,7 @@ class ProcessedResponse:
     handoffs: list[ToolRunHandoff]
     functions: list[ToolRunFunction]
     computer_actions: list[ToolRunComputerAction]
+    tools_used: list[str]  # Names of all tools used, including hosted tools
 
     def has_tools_to_run(self) -> bool:
         # Handoffs, functions and computer actions need local processing
@@ -296,11 +316,24 @@ class RunImpl:
                 next_step=NextStepRunAgain(),
             )
 
+    @classmethod
+    def maybe_reset_tool_choice(
+        cls, agent: Agent[Any], tool_use_tracker: AgentToolUseTracker, model_settings: ModelSettings
+    ) -> ModelSettings:
+        """Resets tool choice to None if the agent has used tools and the agent's reset_tool_choice
+        flag is True."""
+
+        if agent.reset_tool_choice is True and tool_use_tracker.has_used_tools(agent):
+            return dataclasses.replace(model_settings, tool_choice=None)
+
+        return model_settings
+
     @classmethod
     def process_model_response(
         cls,
         *,
         agent: Agent[Any],
+        all_tools: list[Tool],
         response: ModelResponse,
         output_schema: AgentOutputSchema | None,
         handoffs: list[Handoff],
@@ -310,22 +343,25 @@ class RunImpl:
         run_handoffs = []
         functions = []
         computer_actions = []
-
+        tools_used: list[str] = []
         handoff_map = {handoff.tool_name: handoff for handoff in handoffs}
-        function_map = {tool.name: tool for tool in agent.tools if isinstance(tool, FunctionTool)}
-        computer_tool = next((tool for tool in agent.tools if isinstance(tool, ComputerTool)), None)
+        function_map = {tool.name: tool for tool in all_tools if isinstance(tool, FunctionTool)}
+        computer_tool = next((tool for tool in all_tools if isinstance(tool, ComputerTool)), None)
 
         for output in response.output:
             if isinstance(output, ResponseOutputMessage):
                 items.append(MessageOutputItem(raw_item=output, agent=agent))
             elif isinstance(output, ResponseFileSearchToolCall):
                 items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("file_search")
             elif isinstance(output, ResponseFunctionWebSearch):
                 items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("web_search")
             elif isinstance(output, ResponseReasoningItem):
                 items.append(ReasoningItem(raw_item=output, agent=agent))
             elif isinstance(output, ResponseComputerToolCall):
                 items.append(ToolCallItem(raw_item=output, agent=agent))
+                tools_used.append("computer_use")
                 if not computer_tool:
                     _error_tracing.attach_error_to_current_span(
                         SpanError(
@@ -347,6 +383,8 @@ class RunImpl:
             if not isinstance(output, ResponseFunctionToolCall):
                 continue
 
+            tools_used.append(output.name)
+
             # Handoffs
             if output.name in handoff_map:
                 items.append(HandoffCallItem(raw_item=output, agent=agent))
@@ -378,6 +416,7 @@ class RunImpl:
             handoffs=run_handoffs,
             functions=functions,
             computer_actions=computer_actions,
+            tools_used=tools_used,
         )
 
     @classmethod
@@ -490,7 +529,8 @@ class RunImpl:
         run_config: RunConfig,
     ) -> SingleStepResult:
         # If there is more than one handoff, add tool responses that reject those handoffs
-        if len(run_handoffs) > 1:
+        multiple_handoffs = len(run_handoffs) > 1
+        if multiple_handoffs:
             output_message = "Multiple handoffs detected, ignoring this one."
             new_step_items.extend(
                 [
@@ -512,6 +552,16 @@ class RunImpl:
                 context_wrapper, actual_handoff.tool_call.arguments
             )
             span_handoff.span_data.to_agent = new_agent.name
+            if multiple_handoffs:
+                requested_agents = [handoff.handoff.agent_name for handoff in run_handoffs]
+                span_handoff.set_error(
+                    SpanError(
+                        message="Multiple handoffs requested",
+                        data={
+                            "requested_agents": requested_agents,
+                        },
+                    )
+                )
 
             # Append a tool output item for the handoff
             new_step_items.append(

agents/agent.py

@@ -6,12 +6,13 @@ from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, cast
 
-from typing_extensions import TypeAlias, TypedDict
+from typing_extensions import NotRequired, TypeAlias, TypedDict
 
 from .guardrail import InputGuardrail, OutputGuardrail
 from .handoffs import Handoff
 from .items import ItemHelpers
 from .logger import logger
+from .mcp import MCPUtil
 from .model_settings import ModelSettings
 from .models.interface import Model
 from .run_context import RunContextWrapper, TContext
@@ -21,6 +22,7 @@ from .util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
     from .lifecycle import AgentHooks
+    from .mcp import MCPServer
     from .result import RunResult
 
 
@@ -42,7 +44,7 @@ ToolsToFinalOutputFunction: TypeAlias = Callable[
     MaybeAwaitable[ToolsToFinalOutputResult],
 ]
 """A function that takes a run context and a list of tool results, and returns a
-`ToolToFinalOutputResult`.
+`ToolsToFinalOutputResult`.
 """
 
 
@@ -51,6 +53,15 @@ class StopAtTools(TypedDict):
     """A list of tool names, any of which will stop the agent from running further."""
 
 
+class MCPConfig(TypedDict):
+    """Configuration for MCP servers."""
+
+    convert_schemas_to_strict: NotRequired[bool]
+    """If True, we will attempt to convert the MCP schemas to strict-mode schemas. This is a
+    best-effort conversion, so some schemas may not be convertible. Defaults to False.
+    """
+
+
 @dataclass
 class Agent(Generic[TContext]):
     """An agent is an AI model configured with instructions, tools, guardrails, handoffs and more.
@@ -107,6 +118,19 @@ class Agent(Generic[TContext]):
     tools: list[Tool] = field(default_factory=list)
     """A list of tools that the agent can use."""
 
+    mcp_servers: list[MCPServer] = field(default_factory=list)
+    """A list of [Model Context Protocol](https://modelcontextprotocol.io/) servers that
+    the agent can use. Every time the agent runs, it will include tools from these servers in the
+    list of available tools.
+
+    NOTE: You are expected to manage the lifecycle of these servers. Specifically, you must call
+    `server.connect()` before passing it to the agent, and `server.cleanup()` when the server is no
+    longer needed.
+    """
+
+    mcp_config: MCPConfig = field(default_factory=lambda: MCPConfig())
+    """Configuration for MCP servers."""
+
     input_guardrails: list[InputGuardrail[TContext]] = field(default_factory=list)
     """A list of checks that run in parallel to the agent's execution, before generating a
     response. Runs only if the agent is the first agent in the chain.
@@ -143,6 +167,10 @@ class Agent(Generic[TContext]):
       web search, etc are always processed by the LLM.
     """
 
+    reset_tool_choice: bool = True
+    """Whether to reset the tool choice to the default value after a tool has been called. Defaults
+    to True. This ensures that the agent doesn't enter an infinite loop of tool usage."""
+
     def clone(self, **kwargs: Any) -> Agent[TContext]:
         """Make a copy of the agent, with the given arguments changed. For example, you could do:
         ```
@@ -205,3 +233,13 @@ class Agent(Generic[TContext]):
             logger.error(f"Instructions must be a string or a function, got {self.instructions}")
 
         return None
+
+    async def get_mcp_tools(self) -> list[Tool]:
+        """Fetches the available tools from the MCP servers."""
+        convert_schemas_to_strict = self.mcp_config.get("convert_schemas_to_strict", False)
+        return await MCPUtil.get_all_function_tools(self.mcp_servers, convert_schemas_to_strict)
+
+    async def get_all_tools(self) -> list[Tool]:
+        """All agent tools, including MCP tools and function tools."""
+        mcp_tools = await self.get_mcp_tools()
+        return mcp_tools + self.tools

agents/extensions/visualization.py

@@ -0,0 +1,137 @@
+from typing import Optional
+
+import graphviz  # type: ignore
+
+from agents import Agent
+from agents.handoffs import Handoff
+from agents.tool import Tool
+
+
+def get_main_graph(agent: Agent) -> str:
+    """
+    Generates the main graph structure in DOT format for the given agent.
+
+    Args:
+        agent (Agent): The agent for which the graph is to be generated.
+
+    Returns:
+        str: The DOT format string representing the graph.
+    """
+    parts = [
+        """
+    digraph G {
+        graph [splines=true];
+        node [fontname="Arial"];
+        edge [penwidth=1.5];
+    """
+    ]
+    parts.append(get_all_nodes(agent))
+    parts.append(get_all_edges(agent))
+    parts.append("}")
+    return "".join(parts)
+
+
+def get_all_nodes(agent: Agent, parent: Optional[Agent] = None) -> str:
+    """
+    Recursively generates the nodes for the given agent and its handoffs in DOT format.
+
+    Args:
+        agent (Agent): The agent for which the nodes are to be generated.
+
+    Returns:
+        str: The DOT format string representing the nodes.
+    """
+    parts = []
+
+    # Start and end the graph
+    parts.append(
+        '"__start__" [label="__start__", shape=ellipse, style=filled, '
+        "fillcolor=lightblue, width=0.5, height=0.3];"
+        '"__end__" [label="__end__", shape=ellipse, style=filled, '
+        "fillcolor=lightblue, width=0.5, height=0.3];"
+    )
+    # Ensure parent agent node is colored
+    if not parent:
+        parts.append(
+            f'"{agent.name}" [label="{agent.name}", shape=box, style=filled, '
+            "fillcolor=lightyellow, width=1.5, height=0.8];"
+        )
+
+    for tool in agent.tools:
+        parts.append(
+            f'"{tool.name}" [label="{tool.name}", shape=ellipse, style=filled, '
+            f"fillcolor=lightgreen, width=0.5, height=0.3];"
+        )
+
+    for handoff in agent.handoffs:
+        if isinstance(handoff, Handoff):
+            parts.append(
+                f'"{handoff.agent_name}" [label="{handoff.agent_name}", '
+                f"shape=box, style=filled, style=rounded, "
+                f"fillcolor=lightyellow, width=1.5, height=0.8];"
+            )
+        if isinstance(handoff, Agent):
+            parts.append(
+                f'"{handoff.name}" [label="{handoff.name}", '
+                f"shape=box, style=filled, style=rounded, "
+                f"fillcolor=lightyellow, width=1.5, height=0.8];"
+            )
+            parts.append(get_all_nodes(handoff))
+
+    return "".join(parts)
+
+
+def get_all_edges(agent: Agent, parent: Optional[Agent] = None) -> str:
+    """
+    Recursively generates the edges for the given agent and its handoffs in DOT format.
+
+    Args:
+        agent (Agent): The agent for which the edges are to be generated.
+        parent (Agent, optional): The parent agent. Defaults to None.
+
+    Returns:
+        str: The DOT format string representing the edges.
+    """
+    parts = []
+
+    if not parent:
+        parts.append(f'"__start__" -> "{agent.name}";')
+
+    for tool in agent.tools:
+        parts.append(f"""
+        "{agent.name}" -> "{tool.name}" [style=dotted, penwidth=1.5];
+        "{tool.name}" -> "{agent.name}" [style=dotted, penwidth=1.5];""")
+
+    for handoff in agent.handoffs:
+        if isinstance(handoff, Handoff):
+            parts.append(f"""
+            "{agent.name}" -> "{handoff.agent_name}";""")
+        if isinstance(handoff, Agent):
+            parts.append(f"""
+            "{agent.name}" -> "{handoff.name}";""")
+            parts.append(get_all_edges(handoff, agent))
+
+    if not agent.handoffs and not isinstance(agent, Tool):  # type: ignore
+        parts.append(f'"{agent.name}" -> "__end__";')
+
+    return "".join(parts)
+
+
+def draw_graph(agent: Agent, filename: Optional[str] = None) -> graphviz.Source:
+    """
+    Draws the graph for the given agent and optionally saves it as a PNG file.
+
+    Args:
+        agent (Agent): The agent for which the graph is to be drawn.
+        filename (str): The name of the file to save the graph as a PNG.
+
+    Returns:
+        graphviz.Source: The graphviz Source object representing the graph.
+    """
+    dot_code = get_main_graph(agent)
+    graph = graphviz.Source(dot_code)
+
+    if filename:
+        graph.render(filename, format="png")
+
+    return graph

agents/mcp/__init__.py

@@ -0,0 +1,21 @@
+try:
+    from .server import (
+        MCPServer,
+        MCPServerSse,
+        MCPServerSseParams,
+        MCPServerStdio,
+        MCPServerStdioParams,
+    )
+except ImportError:
+    pass
+
+from .util import MCPUtil
+
+__all__ = [
+    "MCPServer",
+    "MCPServerSse",
+    "MCPServerSseParams",
+    "MCPServerStdio",
+    "MCPServerStdioParams",
+    "MCPUtil",
+]