feat: Add Graphviz-based agent visualization functionality

docs/visualization.md

@@ -0,0 +1,106 @@
+# 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
+
+The visualization functionality relies on the **Graphviz** package. To use it, ensure you have Graphviz installed and add it as a dependency in `pyproject.toml`. Alternatively, install it directly via pip:
+
+```bash
+pip install openai-agents[visualization]
+```
+
+## 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.
+
+## Testing the Visualization
+
+The visualization functionality includes test coverage to ensure correctness. Tests are located in `tests/test_visualizations.py` and verify:
+
+- Node and edge correctness in `get_main_graph()`.
+- Proper agent and tool representation in `get_all_nodes()`.
+- Accurate relationship mapping in `get_all_edges()`.
+- Graph rendering functionality in `draw_graph()`.
+
+Run tests using:
+
+```bash
+pytest tests/test_visualizations.py
+```
+
+## Conclusion
+
+Agent visualization provides a powerful way to **understand, debug, and communicate** how agents interact within an application. By leveraging **Graphviz**, you can generate intuitive visual representations of complex agent structures effortlessly.
+
+For further details on agent functionality, see the [Agents documentation](agents.md).
+

pyproject.toml

@@ -49,6 +49,10 @@ dev = [
     "playwright==1.50.0",
     "inline-snapshot>=0.20.7",
 ]
+visualization = [
+    "graphviz>=0.17",
+]
+
 [tool.uv.workspace]
 members = ["agents"]
 

src/agents/extensions/visualization.py

@@ -0,0 +1,119 @@
+import graphviz
+
+from agents import Agent
+
+
+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];
+        "__start__" [shape=ellipse, style=filled, fillcolor=lightblue];
+        "__end__" [shape=ellipse, style=filled, fillcolor=lightblue];
+    """
+    ]
+    parts.append(get_all_nodes(agent))
+    parts.append(get_all_edges(agent))
+    parts.append("}")
+    return "".join(parts)
+
+
+def get_all_nodes(agent: Agent, parent: 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 = []
+    # 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];"
+        )
+
+    # Smaller tools (ellipse, green)
+    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];"
+        )
+
+    # Bigger handoffs (rounded box, yellow)
+    for handoff in agent.handoffs:
+        parts.append(
+            f'"{handoff.name}" [label="{handoff.name}", 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: 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];""")
+
+    if not agent.handoffs:
+        parts.append(f"""
+        "{agent.name}" -> "__end__";""")
+
+    for handoff in agent.handoffs:
+        parts.append(f"""
+        "{agent.name}" -> "{handoff.name}";""")
+        parts.append(get_all_edges(handoff, agent))
+
+    return "".join(parts)
+
+
+def draw_graph(agent: Agent, filename: 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

tests/test_visualization.py

@@ -0,0 +1,116 @@
+from unittest.mock import Mock
+
+import graphviz
+import pytest
+
+from agents import Agent
+from agents.extensions.visualization import (
+    draw_graph,
+    get_all_edges,
+    get_all_nodes,
+    get_main_graph,
+)
+
+
+@pytest.fixture
+def mock_agent():
+    tool1 = Mock()
+    tool1.name = "Tool1"
+    tool2 = Mock()
+    tool2.name = "Tool2"
+
+    handoff1 = Mock()
+    handoff1.name = "Handoff1"
+    handoff1.tools = []
+    handoff1.handoffs = []
+
+    agent = Mock(spec=Agent)
+    agent.name = "Agent1"
+    agent.tools = [tool1, tool2]
+    agent.handoffs = [handoff1]
+
+    return agent
+
+
+def test_get_main_graph(mock_agent):
+    result = get_main_graph(mock_agent)
+    assert "digraph G" in result
+    assert "graph [splines=true];" in result
+    assert 'node [fontname="Arial"];' in result
+    assert "edge [penwidth=1.5];" in result
+    assert '"__start__" [shape=ellipse, style=filled, fillcolor=lightblue];' in result
+    assert '"__end__" [shape=ellipse, style=filled, fillcolor=lightblue];' in result
+    assert (
+        '"Agent1" [label="Agent1", shape=box, style=filled, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in result
+    )
+    assert (
+        '"Tool1" [label="Tool1", shape=ellipse, style=filled, '
+        "fillcolor=lightgreen, width=0.5, height=0.3];" in result
+    )
+    assert (
+        '"Tool2" [label="Tool2", shape=ellipse, style=filled, '
+        "fillcolor=lightgreen, width=0.5, height=0.3];" in result
+    )
+    assert (
+        '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in result
+    )
+
+
+def test_get_all_nodes(mock_agent):
+    result = get_all_nodes(mock_agent)
+    assert (
+        '"Agent1" [label="Agent1", shape=box, style=filled, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in result
+    )
+    assert (
+        '"Tool1" [label="Tool1", shape=ellipse, style=filled, '
+        "fillcolor=lightgreen, width=0.5, height=0.3];" in result
+    )
+    assert (
+        '"Tool2" [label="Tool2", shape=ellipse, style=filled, '
+        "fillcolor=lightgreen, width=0.5, height=0.3];" in result
+    )
+    assert (
+        '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in result
+    )
+
+
+def test_get_all_edges(mock_agent):
+    result = get_all_edges(mock_agent)
+    assert '"__start__" -> "Agent1";' in result
+    assert '"Agent1" -> "Tool1" [style=dotted, penwidth=1.5];' in result
+    assert '"Tool1" -> "Agent1" [style=dotted, penwidth=1.5];' in result
+    assert '"Agent1" -> "Tool2" [style=dotted, penwidth=1.5];' in result
+    assert '"Tool2" -> "Agent1" [style=dotted, penwidth=1.5];' in result
+    assert '"Agent1" -> "Handoff1";' in result
+    assert '"Handoff1" -> "__end__";' in result
+
+
+def test_draw_graph(mock_agent):
+    graph = draw_graph(mock_agent)
+    assert isinstance(graph, graphviz.Source)
+    assert "digraph G" in graph.source
+    assert "graph [splines=true];" in graph.source
+    assert 'node [fontname="Arial"];' in graph.source
+    assert "edge [penwidth=1.5];" in graph.source
+    assert '"__start__" [shape=ellipse, style=filled, fillcolor=lightblue];' in graph.source
+    assert '"__end__" [shape=ellipse, style=filled, fillcolor=lightblue];' in graph.source
+    assert (
+        '"Agent1" [label="Agent1", shape=box, style=filled, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in graph.source
+    )
+    assert (
+        '"Tool1" [label="Tool1", shape=ellipse, style=filled, '
+        "fillcolor=lightgreen, width=0.5, height=0.3];" in graph.source
+    )
+    assert (
+        '"Tool2" [label="Tool2", shape=ellipse, style=filled, '
+        "fillcolor=lightgreen, width=0.5, height=0.3];" in graph.source
+    )
+    assert (
+        '"Handoff1" [label="Handoff1", shape=box, style=filled, style=rounded, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in graph.source
+    )