Merge d248bd8 into sapling-pr-archive-rm-openai

rm-openai · web-flow · commit cd79abf7b47b · 2025-03-26T10:44:08.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -141,4 +141,4 @@ cython_debug/
 .ruff_cache/
 
 # PyPI configuration file
-.pypirc
+.pypirc
diff --git a/docs/assets/images/graph.png b/docs/assets/images/graph.png
diff --git a/docs/visualization.md b/docs/visualization.md
@@ -0,0 +1,86 @@
+# Agent Visualization
+
+Agent visualization allows you to generate a structured graphical representation of agents and their relationships using **Graphviz**. This is useful for understanding how agents, tools, and handoffs interact within an application.
+
+## Installation
+
+Install the optional `viz` dependency group:
+
+```bash
+pip install "openai-agents[viz]"
+```
+
+## Generating a Graph
+
+You can generate an agent visualization using the `draw_graph` function. This function creates a directed graph where:
+
+- **Agents** are represented as yellow boxes.
+- **Tools** are represented as green ellipses.
+- **Handoffs** are directed edges from one agent to another.
+
+### Example Usage
+
+```python
+from agents import Agent, function_tool
+from agents.extensions.visualization import draw_graph
+
+@function_tool
+def get_weather(city: str) -> str:
+    return f"The weather in {city} is sunny."
+
+spanish_agent = Agent(
+    name="Spanish agent",
+    instructions="You only speak Spanish.",
+)
+
+english_agent = Agent(
+    name="English agent",
+    instructions="You only speak English",
+)
+
+triage_agent = Agent(
+    name="Triage agent",
+    instructions="Handoff to the appropriate agent based on the language of the request.",
+    handoffs=[spanish_agent, english_agent],
+    tools=[get_weather],
+)
+
+draw_graph(triage_agent)
+```
+
+![Agent Graph](./assets/images/graph.png)
+
+This generates a graph that visually represents the structure of the **triage agent** and its connections to sub-agents and tools.
+
+
+## Understanding the Visualization
+
+The generated graph includes:
+
+- A **start node** (`__start__`) indicating the entry point.
+- Agents represented as **rectangles** with yellow fill.
+- Tools represented as **ellipses** with green fill.
+- Directed edges indicating interactions:
+  - **Solid arrows** for agent-to-agent handoffs.
+  - **Dotted arrows** for tool invocations.
+- An **end node** (`__end__`) indicating where execution terminates.
+
+## Customizing the Graph
+
+### Showing the Graph
+By default, `draw_graph` displays the graph inline. To show the graph in a separate window, write the following:
+
+```python
+draw_graph(triage_agent).view()
+```
+
+### Saving the Graph
+By default, `draw_graph` displays the graph inline. To save it as a file, specify a filename:
+
+```python
+draw_graph(triage_agent, filename="agent_graph.png")
+```
+
+This will generate `agent_graph.png` in the working directory.
+
+
diff --git a/examples/mcp/filesystem_example/README.md b/examples/mcp/filesystem_example/README.md
@@ -1,6 +1,6 @@
 # MCP Filesystem Example
 
-This example uses the [fileystem MCP server](https://github.com/modelcontextprotocol/servers/tree/main/src/filesystem), running locally via `npx`.
+This example uses the [filesystem MCP server](https://github.com/modelcontextprotocol/servers/tree/main/src/filesystem), running locally via `npx`.
 
 Run it via:
 
diff --git a/examples/mcp/git_example/README.md b/examples/mcp/git_example/README.md
@@ -0,0 +1,25 @@
+# MCP Git Example
+
+This example uses the [git MCP server](https://github.com/modelcontextprotocol/servers/tree/main/src/git), running locally via `uvx`.
+
+Run it via:
+
+```
+uv run python examples/mcp/git_example/main.py
+```
+
+## Details
+
+The example uses the `MCPServerStdio` class from `agents`, with the command:
+
+```bash
+uvx mcp-server-git
+```
+Prior to running the agent, the user is prompted to provide a local directory path to their git repo. Using that, the Agent can invoke Git MCP tools like `git_log` to inspect the git commit log.
+
+Under the hood:
+
+1. The server is spun up in a subprocess, and exposes a bunch of tools like `git_log()`
+2. We add the server instance to the Agent via `mcp_agents`.
+3. Each time the agent runs, we call out to the MCP server to fetch the list of tools via `server.list_tools()`. The result is cached. 
+4. If the LLM chooses to use an MCP tool, we call the MCP server to run the tool via `server.run_tool()`.
diff --git a/examples/mcp/git_example/main.py b/examples/mcp/git_example/main.py
@@ -0,0 +1,48 @@
+import asyncio
+import shutil
+
+from agents import Agent, Runner, trace
+from agents.mcp import MCPServer, MCPServerStdio
+
+
+async def run(mcp_server: MCPServer, directory_path: str):
+    agent = Agent(
+        name="Assistant",
+        instructions=f"Answer questions about the git repository at {directory_path}, use that for repo_path",
+        mcp_servers=[mcp_server],
+    )
+
+    message = "Who's the most frequent contributor?"
+    print("\n" + "-" * 40)
+    print(f"Running: {message}")
+    result = await Runner.run(starting_agent=agent, input=message)
+    print(result.final_output)
+
+    message = "Summarize the last change in the repository."
+    print("\n" + "-" * 40)
+    print(f"Running: {message}")
+    result = await Runner.run(starting_agent=agent, input=message)
+    print(result.final_output)
+
+
+async def main():
+    # Ask the user for the directory path
+    directory_path = input("Please enter the path to the git repository: ")
+
+    async with MCPServerStdio(
+        params={
+            "command": "uvx",
+            "args": [
+                "mcp-server-git"
+            ]
+        }
+    ) as server:
+        with trace(workflow_name="MCP Git Example"):
+            await run(server, directory_path)
+
+
+if __name__ == "__main__":
+    if not shutil.which("uvx"):
+        raise RuntimeError("uvx is not installed. Please install it with `pip install uvx`.")
+
+    asyncio.run(main())
diff --git a/examples/voice/static/main.py b/examples/voice/static/main.py
@@ -1,6 +1,8 @@
 import asyncio
 import random
 
+import numpy as np
+
 from agents import Agent, function_tool
 from agents.extensions.handoff_prompt import prompt_with_handoff_instructions
 from agents.voice import (
@@ -78,6 +80,9 @@ async def main():
             elif event.type == "voice_stream_event_lifecycle":
                 print(f"Received lifecycle event: {event.event}")
 
+        # Add 1 second of silence to the end of the stream to avoid cutting off the last audio.
+        player.add_audio(np.zeros(24000 * 1, dtype=np.int16))
+
 
 if __name__ == "__main__":
     asyncio.run(main())
diff --git a/examples/voice/static/util.py b/examples/voice/static/util.py
@@ -62,6 +62,7 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        self.stream.stop()  # wait for the stream to finish
         self.stream.close()
 
     def add_audio(self, audio_data: npt.NDArray[np.int16]):
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -36,6 +36,7 @@ nav:
       - multi_agent.md
       - models.md
       - config.md
+      - visualization.md
       - Voice agents:
           - voice/quickstart.md
           - voice/pipeline.md
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "openai-agents"
-version = "0.0.6"
+version = "0.0.7"
 description = "OpenAI Agents SDK"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -35,6 +35,7 @@ Repository = "https://github.com/openai/openai-agents-python"
 
 [project.optional-dependencies]
 voice = ["numpy>=2.2.0, <3; python_version>='3.10'", "websockets>=15.0, <16"]
+viz = ["graphviz>=0.17"]
 
 [dependency-groups]
 dev = [
@@ -56,7 +57,9 @@ dev = [
     "pynput",
     "textual",
     "websockets",
+    "graphviz",
 ]
+
 [tool.uv.workspace]
 members = ["agents"]
 
diff --git a/src/agents/extensions/visualization.py b/src/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
diff --git a/tests/test_visualization.py b/tests/test_visualization.py
diff --git a/uv.lock b/uv.lock
