phowell
/
pocketflow
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pocketflow-visualization: adds action names to links

This commit is contained in:
rahoogan 2025-05-16 22:25:45 +10:00
parent d6f36f8c5b
commit c415b553d0
1 changed files with 14 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
cookbook/pocketflow-visualization/visualize.py
View File

@ -28,28 +28,31 @@ def build_mermaid(start):
ids[n] if n in ids else (ids.setdefault(n, f"N{ctr}"), (ctr := ctr + 1))[0] ids[n] if n in ids else (ids.setdefault(n, f"N{ctr}"), (ctr := ctr + 1))[0]
) )
def link(a, b): def link(a, b, action=None):
if action:
lines.append(f" {a} -->|{action}| {b}")
else:
lines.append(f" {a} --> {b}") lines.append(f" {a} --> {b}")
def walk(node, parent=None): def walk(node, parent=None, action=None):
if node in visited: if node in visited:
return parent and link(parent, get_id(node)) return parent and link(parent, get_id(node), action)
visited.add(node) visited.add(node)
if isinstance(node, Flow): if isinstance(node, Flow):
node.start_node and parent and link(parent, get_id(node.start_node)) node.start_node and parent and link(parent, get_id(node.start_node), action)
lines.append( lines.append(
f"\n subgraph sub_flow_{get_id(node)}[{type(node).__name__}]" f"\n subgraph sub_flow_{get_id(node)}[{type(node).__name__}]"
) )
node.start_node and walk(node.start_node) node.start_node and walk(node.start_node)
for nxt in node.successors.values(): for act, nxt in node.successors.items():
node.start_node and walk(nxt, get_id(node.start_node)) or ( node.start_node and walk(nxt, get_id(node.start_node), act) or (
parent and link(parent, get_id(nxt)) parent and link(parent, get_id(nxt), action)
) or walk(nxt) ) or walk(nxt, None, act)
lines.append(" end\n") lines.append(" end\n")
else: else:
lines.append(f" {(nid := get_id(node))}['{type(node).__name__}']") lines.append(f" {(nid := get_id(node))}['{type(node).__name__}']")
parent and link(parent, nid) parent and link(parent, nid, action)
[walk(nxt, nid) for nxt in node.successors.values()] [walk(nxt, nid, act) for act, nxt in node.successors.items()]
walk(start) walk(start)
return "\n".join(lines) return "\n".join(lines)