Skip to content

Commit 5f96734

Browse files
maximilienmaximilien
authored
first version of Workflow::to_mermaid method (#296)
* first version of Workflow::to_mermaid method * fixed /step.get['agent']/step.get('agent')/s and also updated mermaid flowchart and sequenceDiagram links * import unittest
1 parent 3058dff commit 5f96734

File tree

4 files changed

+200
-1
lines changed

4 files changed

+200
-1
lines changed

‎maestro/src/mermaid.py

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
#!/usr/bin/env python3
2+
3+
# Copyright © 2025 IBM
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
class Mermaid:
18+
# kind: sequenceDiagram or flowchart
19+
# orientation: TD (top down),
20+
# RL (right left)
21+
# when kind is flowchart
22+
def __init__(self, workflow, kind="sequenceDiagram", orientation="TD"):
23+
self.workflow = workflow[0]
24+
self.kind = kind
25+
self.orientation = orientation
26+
27+
# generates a mermaid markdown representation of the workflow
28+
def to_markdown(self) -> str:
29+
sb, markdown = "", ""
30+
if self.kind == "sequenceDiagram":
31+
markdown = self.__to_sequenceDiagram(sb)
32+
elif self.kind == "flowchart":
33+
markdown = self.__to_flowchart(sb, self.orientation)
34+
else:
35+
raise RuntimeError(f"Invalid Mermaid kind: {kind}")
36+
return markdown
37+
38+
# private methods
39+
40+
# returns a markdown of the workflow as a mermaid sequence diagram
41+
#
42+
# sequenceDiagram
43+
# participant agent1
44+
# participant agent2
45+
#
46+
# agent1->>agent2: step1
47+
# agent2->>agent3: step2
48+
# agent2-->>agent1: step3
49+
# agent1->>agent3: step4
50+
#
51+
# See mermaid sequence diagram documentation:
52+
# https://mermaid.js.org/syntax/sequenceDiagram.html
53+
def __to_sequenceDiagram(self, sb):
54+
sb += "sequenceDiagram\n"
55+
for agent in self.workflow['spec']['template']['agents']:
56+
sb += f"participant {agent}\n"
57+
steps, i = self.workflow['spec']['template']['steps'], 0
58+
for step in steps:
59+
agentL = step.get('agent')
60+
agentR = None
61+
if i < (len(steps) - 1):
62+
agentR = steps[i+1].get('agent')
63+
if agentR != None:
64+
sb += f"{agentL}->>{agentR}: {step['name']}\n"
65+
else:
66+
sb += f"{agentL}->>{agentL}: {step['name']}\n"
67+
i = i + 1
68+
return sb
69+
70+
# returns a markdown of the workflow as a mermaid sequence diagram
71+
#
72+
# flowchart LR
73+
# agemt1-- step1 -->agent2
74+
# agemt2-- step2 -->agent3
75+
# agemt3-- step3 -->agent3
76+
#
77+
# See mermaid sequence diagram documentation:
78+
# https://mermaid.js.org/syntax/flowchart.html
79+
def __to_flowchart(self, sb, orientation):
80+
sb += f"flowchart {orientation}\n"
81+
steps, i = self.workflow['spec']['template']['steps'], 0
82+
for step in steps:
83+
agentL = step.get('agent')
84+
agentR = None
85+
if i < (len(steps) - 1):
86+
agentR = steps[i+1].get('agent')
87+
if agentR != None:
88+
sb += f"{agentL}-- {step['name']} -->{agentR}\n"
89+
else:
90+
sb += f"{agentL}-- {step['name']} -->{agentL}\n"
91+
i = i + 1
92+
return sb
93+
94+

‎maestro/src/workflow.py

+8
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
import os, dotenv
1818

19+
from src.mermaid import Mermaid
20+
1921
from src.step import Step
2022
from src.agents.agent_factory import AgentFramework
2123

@@ -56,6 +58,12 @@ def __init__(self, agent_defs={}, workflow={}):
5658
self.agent_defs = agent_defs
5759
self.workflow = workflow
5860

61+
# generates a mermaid markdown representation of the workflow
62+
# kind: sequenceDiagram or flowchart
63+
# orientation: TD (top down) or RL (right left) when kind is flowchart
64+
def to_mermaid(self, kind="sequenceDiagram", orientation="TD") -> str:
65+
return Mermaid(self.workflow, kind, orientation).to_markdown()
66+
5967
async def run(self):
6068
"""Execute workflow."""
6169
self.create_or_restore_agents(self.agent_defs, self.workflow)

‎maestro/tests/workflow/test_mermaid.py

+97
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
#!/usr/bin/env python3
2+
3+
# Copyright © 2025 IBM
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
import os, yaml, unittest
18+
19+
from unittest import TestCase
20+
21+
from src.mermaid import Mermaid
22+
from src.workflow import Workflow
23+
24+
def parse_yaml(file_path):
25+
with open(file_path, "r") as file:
26+
yaml_data = list(yaml.safe_load_all(file))
27+
return yaml_data
28+
29+
class TestMermaid(TestCase):
30+
def setUp(self):
31+
self.workflow_yaml = parse_yaml(os.path.join(os.path.dirname(__file__),"../yamls/workflows/simple_workflow.yaml"))
32+
33+
def tearDown(self):
34+
self.workflow_yaml = None
35+
36+
def test_markdown_sequenceDiagram(self):
37+
mermaid = Mermaid(self.workflow_yaml, "sequenceDiagram")
38+
self.assertTrue(mermaid.to_markdown().startswith("sequenceDiagram"))
39+
40+
for agent in ['test1', 'test2', 'test3']:
41+
self.assertTrue(f"participant {agent}" in mermaid.to_markdown())
42+
43+
self.assertTrue(f"test1->>test2: step1" in mermaid.to_markdown())
44+
self.assertTrue(f"test2->>test3: step2" in mermaid.to_markdown())
45+
self.assertTrue(f"test3->>test3: step3" in mermaid.to_markdown())
46+
47+
def test_markdown_flowchartTD(self):
48+
mermaid = Mermaid(self.workflow_yaml, "flowchart", "TD")
49+
self.assertTrue(mermaid.to_markdown().startswith("flowchart TD"))
50+
51+
self.assertTrue(f"test1-- step1 -->test2" in mermaid.to_markdown())
52+
self.assertTrue(f"test2-- step2 -->test3" in mermaid.to_markdown())
53+
self.assertTrue(f"test3-- step3 -->test3" in mermaid.to_markdown())
54+
55+
def test_markdown_flowchartRL(self):
56+
mermaid = Mermaid(self.workflow_yaml, "flowchart", "RL")
57+
self.assertTrue(mermaid.to_markdown().startswith("flowchart RL"))
58+
59+
self.assertTrue(f"test1-- step1 -->test2" in mermaid.to_markdown())
60+
self.assertTrue(f"test2-- step2 -->test3" in mermaid.to_markdown())
61+
self.assertTrue(f"test3-- step3 -->test3" in mermaid.to_markdown())
62+
63+
class TestWorkdlow_to_mermaid(TestCase):
64+
def setUp(self):
65+
self.workflow_yaml = parse_yaml(os.path.join(os.path.dirname(__file__),"../yamls/workflows/simple_workflow.yaml"))
66+
self.workflow = Workflow({}, self.workflow_yaml)
67+
68+
def tearDown(self):
69+
self.workflow_yaml = None
70+
self.workflow = None
71+
72+
def test_markdown_sequenceDiagram(self):
73+
markdown = self.workflow.to_mermaid("sequenceDiagram")
74+
self.assertTrue(markdown.startswith("sequenceDiagram"))
75+
76+
self.assertTrue(f"test1->>test2: step1" in markdown)
77+
self.assertTrue(f"test2->>test3: step2" in markdown)
78+
self.assertTrue(f"test3->>test3: step3" in markdown)
79+
80+
def test_markdown_flowchartTD(self):
81+
markdown = self.workflow.to_mermaid("flowchart", "TD")
82+
self.assertTrue(markdown.startswith("flowchart TD"))
83+
84+
self.assertTrue(f"test1-- step1 -->test2" in markdown)
85+
self.assertTrue(f"test2-- step2 -->test3" in markdown)
86+
self.assertTrue(f"test3-- step3 -->test3" in markdown)
87+
88+
def test_markdown_flowchartRL(self):
89+
markdown = self.workflow.to_mermaid("flowchart", "RL")
90+
self.assertTrue(markdown.startswith("flowchart RL"))
91+
92+
self.assertTrue(f"test1-- step1 -->test2" in markdown)
93+
self.assertTrue(f"test2-- step2 -->test3" in markdown)
94+
self.assertTrue(f"test3-- step3 -->test3" in markdown)
95+
96+
if __name__ == '__main__':
97+
unittest.main()

‎maestro/tests/yamls/workflows/simple_workflow.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ spec:
1616
- test2
1717
- test3
1818
prompt: This is a test input
19-
exception: step4
19+
exception: step3
2020
steps:
2121
- name: step1
2222
agent: test1

0 commit comments

Comments
 (0)