8000 visualize agent workflows through mermaid · google/adk-python@17da718 · GitHub
[go: up one dir, main page]

Skip to content

Commit 17da718

Browse files
committed
visualize agent workflows through mermaid
1 parent be71208 commit 17da718

File tree

2 files changed

+250
-0
lines changed

2 files changed

+250
-0
lines changed

src/google/adk/utils/visualization.py

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
"""
2+
Utilities for visualizing google-adk agents.
3+
"""
4+
5+
from __future__ import annotations
6+
7+
import itertools
8+
from typing import Any
9+
from typing import Tuple
10+
11+
from google.adk.agents import LoopAgent
12+
from google.adk.agents import ParallelAgent
13+
from google.adk.agents import SequentialAgent
14+
import requests
15+
16+
17+
def build_mermaid(root_agent: Any) -> Tuple[str, bytes]:
18+
"""
19+
Generates a Mermaid 'flowchart LR' diagram for a google-adk
20+
agent tree and returns both the Mermaid source and a PNG
21+
image rendered via the Kroki API.
22+
23+
Args:
24+
root_agent (Any):
25+
The root agent node of the google-adk agent tree.
26+
This should be an instance
27+
of SequentialAgent, LoopAgent, ParallelAgent,
28+
or a compatible agent class with a
29+
`name` attribute and an optional `sub_agents`
30+
attribute.
31+
32+
Returns:
33+
Tuple[str, bytes]:
34+
A tuple containing:
35+
- The Mermaid source code as a string.
36+
- The PNG image bytes rendered from the Mermaid diagram.
37+
38+
Raises:
39+
requests.RequestException: If the request to the Kroki API fails.
40+
41+
Example:
42+
>>> mermaid_src, png_bytes = build_mermaid(my_agent_tree)
43+
>>> print(mermaid_src)
44+
>>> with open("diagram.png", "wb") as f:
45+
... f.write(png_bytes)
46+
"""
47+
clusters, edges = [], []
48+
first_of, last_of, nodes = {}, {}, {}
49+
50+
# Walk the agent tree
51+
def walk(node):
52+
nid = id(node)
53+
nodes[nid] = node
54+
name = node.name
55+
subs = getattr(node, "sub_agents", []) or []
56+
if subs:
57+
first_of[nid], last_of[nid] = subs[0].name, subs[-1].name
58+
# Create subgraph for non-root composite nodes
59+
if node is not root_agent and isinstance(
60+
node, (SequentialAgent, LoopAgent, ParallelAgent)
61+
):
62+
block = [f'subgraph {name}["{name}"]']
63+
if isinstance(node, (SequentialAgent, LoopAgent)):
64+
for a, b in itertools.pairwise(subs):
65+
block.append(f" {a.name} --> {b.name}")
66+
# loop-back even for single-child loops
67+
if isinstance(node, LoopAgent):
68+
if len(subs) == 1:
69+
block.append(f" {subs[0].name} -.->|repeat| {subs[0].name}")
70+
elif len(subs) > 1:
71+
block.append(f" {subs[-1].name} -.->|repeat| {subs[0].name}")
72+
elif isinstance(node, ParallelAgent):
73+
for child in subs:
74+
block.append(f' {child.name}["{child.name}"]')
75+
block.append("end")
76+
clusters.append("\n".join(block))
77+
# Recurse
78+
for child in subs:
79+
walk(child)
80+
81+
walk(root_agent)
82+
83+
# Link root children
84+
if isinstance(root_agent, SequentialAgent):
85+
children = root_agent.sub_agents or []
86+
# Kick-off
87+
if children:
88+
first = children[0]
89+
if isinstance(first, ParallelAgent):
90+
for c in first.sub_agents:
91+
edges.append(f"{root_agent.name} -.-> {c.name}")
92+
else:
93+
edges.append(
94+
f"{root_agent.name} -.-> {first_of.get(id(first), first.name)}"
95+
)
96+
# Chain
97+
for prev, nxt in itertools.pairwise(children):
98+
prev_exits = (
99+
[c.name for c in prev.sub_agents]
100+
if isinstance(prev, ParallelAgent)
101+
else [last_of.get(id(prev), prev.name)]
102+
)
103+
nxt_entries = (
104+
[c.name for c in nxt.sub_agents]
105+
if isinstance(nxt, ParallelAgent)
106+
else [first_of.get(id(nxt), nxt.name)]
107+
)
108+
arrow = "-.->" if isinstance(nxt, ParallelAgent) else "-->"
109+
for src in prev_exits:
110+
for dst in nxt_entries:
111+
edges.append(f"{src} {arrow} {dst}")
112+
else:
113+
for c in getattr(root_agent, "sub_agents", []) or []:
114+
edges.append(f"{root_agent.name} --> {c.name}")
115+
116+
# Assemble graph as mermaid code
117+
mermaid_src = "\n".join(
118+
["flowchart LR", f'{root_agent.name}["{root_agent.name}"]']
119+
+ clusters
120+
+ edges
121+
)
122+
123+
# Render via Kroki
124+
# note: kroki is a third party service which enables the rendering
125+
# of mermaid diagrams without local npm installation of mermaid.
126+
png = requests.post(
127+
"https://kroki.io/mermaid/png",
128+
data=mermaid_src.encode("utf-8"),
129+
headers={"Content-Type": "text/plain"},
130+
).content
131+
132+
return mermaid_src, png
Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
"""
2+
Tests for src/google/adk/utils/visualization.py
3+
"""
4+
5+
from __future__ import annotations
6+
7+
from google.adk.agents import Agent
8+
from google.adk.agents import LoopAgent
9+
from google.adk.agents import ParallelAgent
10+
from google.adk.agents import SequentialAgent
11+
from google.adk.utils.visualization import build_mermaid
12+
13+
14+
def test_build_mermaid():
15+
"""
16+
Tests build_mermaid function.
17+
18+
We build an agent workflow and then pass the
19+
root_agent to build_mermaid, which builds
20+
a mermaid diagram.
21+
"""
22+
23+
agent1 = Agent(
24+
model="nonexistent-model",
25+
name="agent1",
26+
description="Example",
27+
instruction="""
28+
Example
29+
""",
30+
)
31+
< 10000 /code>
32+
agent2 = Agent(
33+
model="nonexistent-model",
34+
name="agent2",
35+
description="Example",
36+
instruction="""
37+
Example
38+
""",
39+
)
40+
41+
agent3 = Agent(
42+
model="nonexistent-model",
43+
name="agent3",
44+
description="Example",
45+
instruction=f"""
46+
Example
47+
""",
48+
)
49+
50+
agent4 = Agent(
51+
model="nonexistent-model",
52+
name="agent4",
53+
description="Example",
54+
instruction=f"""
55+
Example
56+
""",
57+
)
58+
59+
agent5 = Agent(
60+
model="nonexistent-model",
61+
name="agent5",
62+
description="Example",
63+
instruction=f"""
64+
Example
65+
""",
66+
)
67+
68+
agent6 = Agent(
69+
model="nonexistent-model",
70+
name="agent6",
71+
description="Example",
72+
instruction=f"""
73+
Example
74+
""",
75+
)
76+
77+
agent7 = Agent(
78+
model="nonexistent-model",
79+
name="agent7",
80+
description="Example",
81+
instruction=f"""
82+
Example
83+
""",
84+
)
85+
86+
# example sequence
87+
sequence_1 = SequentialAgent(
88+
name="ExampleSequence",
89+
sub_agents=[agent1, agent2],
90+
)
91+
92+
# example loop
93+
loop_1 = LoopAgent(
94+
name="ExampleLoop",
95+
sub_agents=[agent6, agent7],
96+
max_iterations=10,
97+
)
98+
99+
# example parallel
100+
parallel_1 = ParallelAgent(
101+
name="ExampleParallel",
102+
sub_agents=[agent3, agent4, agent5],
103+
)
104+
105+
# sequence for orchestrating everything together
106+
root_agent = SequentialAgent(
107+
name="root_agent",
108+
sub_agents=[sequence_1, loop_1, parallel_1],
109+
description="Example",
110+
)
111+
112+
mermaid_src, png_display_bytes = build_mermaid(root_agent)
113+
114+
assert isinstance(mermaid_src, str)
115+
assert mermaid_src.startswith("flowchart LR")
116+
117+
assert isinstance(png_display_bytes, bytes)
118+
assert png_display_bytes.startswith(b"\x89PNG\r\n\x1a\n")

0 commit comments

Comments
 (0)
0