Spaces:
Running
Running
File size: 4,699 Bytes
d631808 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
from __future__ import annotations
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: Agent | None = None, visited: set[str] | None = 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.
"""
if visited is None:
visited = set()
if agent.name in visited:
return ""
visited.add(agent.name)
parts = []
# Start and end the graph
if not parent:
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
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):
if handoff.name not in visited:
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, agent, visited))
return "".join(parts)
def get_all_edges(
agent: Agent, parent: Agent | None = None, visited: set[str] | None = 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.
"""
if visited is None:
visited = set()
if agent.name in visited:
return ""
visited.add(agent.name)
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, visited))
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: str | None = 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", cleanup=True)
return graph
|