Plots the workflow graph with node statuses.
(
app_info: dict[str, Any],
agent_state: dict[str, Any] | None = None,
format: str = "svg",
dark_mode: bool = True,
)
| 25 | |
| 26 | |
| 27 | def plot_workflow_graph( |
| 28 | app_info: dict[str, Any], |
| 29 | agent_state: dict[str, Any] | None = None, |
| 30 | format: str = "svg", |
| 31 | dark_mode: bool = True, |
| 32 | ) -> str | bytes: |
| 33 | """Plots the workflow graph with node statuses.""" |
| 34 | agent_state = agent_state or {} |
| 35 | root_agent = app_info.get("root_agent", {}) |
| 36 | graph = root_agent.get("graph", {}) |
| 37 | is_workflow = bool(graph) |
| 38 | |
| 39 | if not graph: |
| 40 | root_name = root_agent.get("name", "root_agent") |
| 41 | sub_agents = root_agent.get("sub_agents", []) |
| 42 | tools = root_agent.get("tools", []) |
| 43 | |
| 44 | nodes = [{"name": root_name, "type": "agent", "tools": tools}] |
| 45 | edges = [] |
| 46 | |
| 47 | def _traverse_sub_agents(agent_dict, parent_name): |
| 48 | for sub in agent_dict.get("sub_agents", []): |
| 49 | sub_name = sub.get("name") |
| 50 | if sub_name: |
| 51 | nodes.append( |
| 52 | {"name": sub_name, "type": "agent", "tools": sub.get("tools", [])} |
| 53 | ) |
| 54 | edges.append({ |
| 55 | "from_node": {"name": parent_name}, |
| 56 | "to_node": {"name": sub_name}, |
| 57 | }) |
| 58 | _traverse_sub_agents(sub, sub_name) |
| 59 | |
| 60 | _traverse_sub_agents(root_agent, root_name) |
| 61 | graph = {"nodes": nodes, "edges": edges} |
| 62 | |
| 63 | nodes_state = agent_state.get("nodes", {}) |
| 64 | dot = graphviz.Digraph(comment="Workflow Visualization") |
| 65 | |
| 66 | if dark_mode: |
| 67 | graph_bgcolor = "#0F172A" |
| 68 | node_fillcolor = "#1E293B" |
| 69 | node_color = "#475569" |
| 70 | node_fontcolor = "#F8FAFC" |
| 71 | edge_color = "#94A3B8" |
| 72 | edge_fontcolor = "#CBD5E1" |
| 73 | start_fillcolor = "#059669" |
| 74 | start_color = "#047857" |
| 75 | end_fillcolor = "#DC2626" |
| 76 | end_color = "#B91C1C" |
| 77 | status_colors = { |
| 78 | NodeStatus.COMPLETED: "#16A34A", |
| 79 | NodeStatus.RUNNING: "#D97706", |
| 80 | NodeStatus.FAILED: "#EF4444", |
| 81 | NodeStatus.INACTIVE: "#1E293B", |
| 82 | NodeStatus.WAITING: "#9333EA", |
| 83 | NodeStatus.CANCELLED: "#475569", |
| 84 | } |
no test coverage detected