run-llama/workflows-py
1
0
Fork 0
mirror of https://github.com/run-llama/workflows-py.git synced 2026-08-24 20:01:34 -04:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

Add draw_most_recent_execution_mermaid (#265)

Browse Source 
* feat: Add draw_most_recent_execution_mermaid function

Co-authored-by: cecli (gemini/gemini-2.5-pro)

* Draw Mermaid: Refactored code.  Not finished.

Draw Mermaid: More refactor

Draw Mermaid: More Refactor to clean up

Draw Mermaid: More refactor.

Draw Mermaid: Nicely Refactored

* Added mermaid code for draw_most_recent_execution_mermaid and test

test: Add tests for draw_most_recent_execution

Co-authored-by: cecli (gemini/gemini-2.5-pro)

Draw Mermaid: Added in the test case for draw_recent_execution_mermaid

* test case for draw_most_recent_execution_mermaid

* Moved import json to the top

* moved import json from second line to first line

* fix sort and bump lock

---------

Co-authored-by: Your Name <you@example.com>
Co-authored-by: Adrian Lyjak <adrianlyjak@gmail.com>
This commit is contained in:
hippopond
2026-01-07 13:41:33 -08:00
committed by GitHub
parent 1544647770
commit 96fd9c9bfb
6 changed files with 219 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/draw_recent_execution_mermaid.md
+7
View File
@@ -0,0 +1,7 @@
---
"llama-index-utils-workflows": minor
---
Added new function draw_most_recent_execution_mermaid
To draw the most recent workflow run in a mermaid format
.gitignore
+1
View File
@@ -149,3 +149,4 @@ cython_debug/
openapi.json
workflow_all_flows.mermaid
node_modules/
.cecli*
packages/llama-index-utils-workflow/pyproject.toml
+6 -1
View File
@@ -3,7 +3,12 @@ requires = ["uv_build>=0.9.10,<0.10.0"]
build-backend = "uv_build"
[dependency-groups]
dev = ["pytest>=8.4.2", "pytest-asyncio>=1.0.0", "pytest-cov>=6.1.1"]
dev = [
"pre-commit>=4.3.0",
"pytest>=8.4.2",
"pytest-asyncio>=1.0.0",
"pytest-cov>=6.1.1"
]
[project]
name = "llama-index-utils-workflow"
packages/llama-index-utils-workflow/src/llama_index/utils/workflow/__init__.py
+168 -113
View File
@@ -1,8 +1,8 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
from __future__ import annotations
import json
from typing import Any, Dict, List, Tuple, Union, cast
from llama_index.core.agent.workflow import (
@@ -410,6 +410,70 @@ def _extract_agent_workflow_structure(
return DrawWorkflowGraph(nodes=nodes, edges=edges)
def _extract_execution_graph(
handler: WorkflowHandler, max_label_length: int | None = None
) -> Tuple[Dict[str, Tuple[str, str, type | None]], List[Tuple[str, str]]]:
"""Helper to extract nodes and edges from the workflow handler's tick log."""
if handler.ctx is None or handler.ctx._broker_run is None:
raise ValueError("No context/run info in this handler. Has it been run yet?")
ticks: List[WorkflowTick] = handler.ctx._broker_run._tick_log
nodes: Dict[str, Tuple[str, str, type | None]] = {}
edges: List[Tuple[str, str]] = []
event_node_by_identity: Dict[int, str] = {}
step_seq: Dict[str, int] = {}
external_node_id = "external_step"
nodes[external_node_id] = ("external_step", "external", None)
def ensure_event_node(ev: Event) -> str:
key = id(ev)
if key in event_node_by_identity:
return event_node_by_identity[key]
label = type(ev).__name__
node_id = f"event:{label}#{len(event_node_by_identity)}"
display_label = (
_truncate_label(label, max_label_length) if max_label_length else label
)
nodes[node_id] = (display_label, "event", type(ev))
event_node_by_identity[key] = node_id
return node_id
def iter_emitted_events(step_tick: TickStepResult[Any]) -> List[Event]:
emitted: List[Event] = []
for r in step_tick.result:
if isinstance(r, StepWorkerResult) and isinstance(r.result, Event):
emitted.append(r.result)
elif isinstance(r, AddCollectedEvent):
emitted.append(r.event)
return emitted
for t in ticks:
if isinstance(t, TickAddEvent):
ev_id = ensure_event_node(t.event)
edges.append((external_node_id, ev_id))
elif isinstance(t, TickStepResult):
step_seq[t.step_name] = step_seq.get(t.step_name, 0) + 1
seq = step_seq[t.step_name]
step_node_id = f"step:{t.step_name}#{seq}"
step_label = f"{t.step_name}#{seq}"
display_label = (
_truncate_label(step_label, max_label_length)
if max_label_length
else step_label
)
nodes[step_node_id] = (display_label, "step", None)
in_event_node_id = ensure_event_node(t.event)
edges.append((in_event_node_id, step_node_id))
for out_ev in iter_emitted_events(t):
out_event_node_id = ensure_event_node(out_ev)
edges.append((step_node_id, out_event_node_id))
return nodes, edges
def draw_all_possible_flows(
workflow: Workflow,
filename: str = "workflow_all_flows.html",
@@ -430,118 +494,6 @@ def draw_all_possible_flows(
_render_pyvis(graph, filename, notebook)
def draw_most_recent_execution(
handler: WorkflowHandler,
filename: str = "workflow_recent_execution.html",
notebook: bool = False,
max_label_length: int | None = None,
) -> None:
"""
Draws the most recent execution of the workflow.
Args:
workflow: The workflow to visualize
filename: Output HTML filename
notebook: Whether running in notebook environment
max_label_length: Maximum label length before truncation (None = no limit)
"""
net = Network(directed=True, height="750px", width="100%")
if handler.ctx is None or handler.ctx._broker_run is None:
raise ValueError("No context/run info in this handler. Has it been run yet?")
ticks: List[WorkflowTick] = handler.ctx._broker_run._tick_log
# Build execution DAG from ticks
nodes: Dict[str, Tuple[str, str, type | None]] = {}
edges: List[Tuple[str, str]] = []
event_node_by_identity: Dict[int, str] = {}
step_seq: Dict[str, int] = {}
# Optional external node for externally added events
external_node_id = "external_step"
nodes[external_node_id] = ("external_step", "external", None)
def ensure_event_node(ev: Event) -> str:
key = id(ev)
if key in event_node_by_identity:
return event_node_by_identity[key]
label = type(ev).__name__
node_id = f"event:{label}#{len(event_node_by_identity)}"
# Truncate label if requested (node label only, id remains stable)
display_label = (
_truncate_label(label, max_label_length) if max_label_length else label
)
nodes[node_id] = (display_label, "event", type(ev))
event_node_by_identity[key] = node_id
return node_id
# Helper to enumerate events emitted from a step result list
def iter_emitted_events(step_tick: TickStepResult[Any]) -> List[Event]:
emitted: List[Event] = []
for r in step_tick.result:
if isinstance(r, StepWorkerResult) and isinstance(r.result, Event):
emitted.append(r.result)
elif isinstance(r, AddCollectedEvent):
emitted.append(r.event)
return emitted
for idx, t in enumerate(ticks):
if isinstance(t, TickAddEvent):
ev_id = ensure_event_node(t.event)
edges.append((external_node_id, ev_id))
elif isinstance(t, TickStepResult):
# Create a step execution node
step_seq[t.step_name] = step_seq.get(t.step_name, 0) + 1
seq = step_seq[t.step_name]
step_node_id = f"step:{t.step_name}#{seq}"
step_label = f"{t.step_name}#{seq}"
display_label = (
_truncate_label(step_label, max_label_length)
if max_label_length
else step_label
)
nodes[step_node_id] = (display_label, "step", None)
# consumed event -> step
in_event_node_id = ensure_event_node(t.event)
edges.append((in_event_node_id, step_node_id))
# step -> emitted events
for out_ev in iter_emitted_events(t):
out_event_node_id = ensure_event_node(out_ev)
edges.append((step_node_id, out_event_node_id))
# Render with Pyvis
# Add nodes first
for node_id, (label, node_type, ev_type) in nodes.items():
if node_type == "step":
color = "#ADD8E6"
shape = "box"
elif node_type == "external":
color = "#BEDAE4"
shape = "box"
else:
# event
color = _determine_event_color(ev_type if ev_type else Event)
shape = "ellipse"
net.add_node(node_id, label=label, color=color, shape=shape)
# Then edges
for src, dst in edges:
net.add_edge(src, dst)
# Suggest a hierarchical layout to preserve timeOrder-ish viewing
try:
net.set_options(
'{"layout": {"hierarchical": {"enabled": true, "direction": "LR", "nodeSpacing": 150, "levelSeparation": 120}}, "physics": {"enabled": false}}'
)
except Exception:
pass
net.show(filename, notebook=notebook)
def draw_all_possible_flows_mermaid(
workflow: Workflow,
filename: str = "workflow_all_flows.mermaid",
@@ -645,3 +597,106 @@ def draw_agent_workflow_mermaid(
"""
graph = _extract_agent_workflow_structure(agent_workflow)
return _render_mermaid(graph, filename)
def draw_most_recent_execution(
handler: WorkflowHandler,
filename: str = "workflow_recent_execution.html",
notebook: bool = False,
max_label_length: int | None = None,
) -> None:
"""Draws the most recent execution of the workflow using Pyvis."""
nodes, edges = _extract_execution_graph(handler, max_label_length)
net = Network(directed=True, height="750px", width="100%")
for node_id, (label, node_type, ev_type) in nodes.items():
if node_type == "step" or node_type == "external":
color = "#ADD8E6" if node_type == "step" else "#BEDAE4"
shape = "box"
else:
color = _determine_event_color(ev_type if ev_type else Event)
shape = "ellipse"
net.add_node(node_id, label=label, color=color, shape=shape)
for src, dst in edges:
net.add_edge(src, dst)
options = {
"layout": {
"hierarchical": {
"enabled": True,
"direction": "LR",
"nodeSpacing": 150,
"levelSeparation": 120,
}
},
"physics": {"enabled": False},
}
try:
net.set_options(json.dumps(options))
except Exception:
pass
net.show(filename, notebook=notebook)
def draw_most_recent_execution_mermaid(
handler: WorkflowHandler,
filename: str = "workflow_recent_execution.mermaid",
max_label_length: int | None = None,
) -> str:
"""Draws the most recent execution of the workflow as a Mermaid diagram."""
nodes, edges = _extract_execution_graph(handler, max_label_length)
mermaid_lines = ["flowchart TD"]
cleaned_ids = {
node_id: node_id.replace(":", "_").replace("#", "_") for node_id in nodes.keys()
}
for node_id, (label, node_type, ev_type) in nodes.items():
clean_id = cleaned_ids[node_id]
shape_start, shape_end = (
("[", "]") if node_type in ["step", "external"] else ("([", "])")
)
css_class = "defaultEventStyle"
if node_type == "step":
css_class = "stepStyle"
elif node_type == "external":
css_class = "externalStyle"
elif node_type == "event" and ev_type:
if issubclass(ev_type, StartEvent):
css_class = "startEventStyle"
elif issubclass(ev_type, StopEvent):
css_class = "stopEventStyle"
mermaid_lines.append(
f' {clean_id}{shape_start}"{label}"{shape_end}:::{css_class}'
)
for src, dst in edges:
mermaid_lines.append(f" {cleaned_ids[src]} --> {cleaned_ids[dst]}")
styles = [
"classDef stepStyle fill:#ADD8E6,color:#000000,line-height:1.2",
"classDef externalStyle fill:#BEDAE4,color:#000000,line-height:1.2",
"classDef startEventStyle fill:#E27AFF,color:#000000",
"classDef stopEventStyle fill:#FFA07A,color:#000000",
"classDef defaultEventStyle fill:#90EE90,color:#000000",
"classDef reactAgentStyle fill:#E27AFF,color:#000000",
"classDef codeActAgentStyle fill:#66ccff,color:#000000",
"classDef defaultAgentStyle fill:#90EE90,color:#000000",
"classDef toolStyle fill:#ff9966,color:#000000",
"classDef workflowBaseStyle fill:#90EE90,color:#000000",
"classDef workflowAgentStyle fill:#66ccff,color:#000000",
"classDef workflowToolStyle fill:#ff9966,color:#000000",
"classDef workflowHandoffStyle fill:#E27AFF,color:#000000",
]
mermaid_lines.extend([f" {s}" for s in styles])
diagram_string = "\n".join(mermaid_lines)
if filename:
with open(filename, "w") as f:
f.write(diagram_string)
return diagram_string
packages/llama-index-utils-workflow/tests/test_drawing.py
+34
View File
@@ -5,6 +5,7 @@ from llama_index.utils.workflow import (
draw_all_possible_flows,
draw_all_possible_flows_mermaid,
draw_most_recent_execution,
draw_most_recent_execution_mermaid,
)
from workflows.workflow import Workflow
@@ -275,3 +276,36 @@ async def test_mermaid_empty_filename(workflow: Workflow) -> None:
# Both should be identical
assert result1 == result2
@pytest.mark.asyncio
async def test_draw_most_recent_execution_mermaid(workflow: Workflow) -> None:
"""Test Mermaid diagram generation for the most recent execution."""
handler = workflow.run()
await handler
with patch("builtins.open", mock_open()) as mock_file:
result = draw_most_recent_execution_mermaid(
handler, filename="test_recent.mermaid"
)
# Verify file was written
mock_file.assert_called_once_with("test_recent.mermaid", "w")
# Verify basic structure
assert isinstance(result, str)
assert result.startswith("flowchart TD")
# Verify it contains style definitions
assert "classDef stepStyle fill:#ADD8E6" in result
assert "classDef startEventStyle fill:#E27AFF" in result
assert "classDef stopEventStyle fill:#FFA07A" in result
assert "classDef defaultEventStyle fill:#90EE90" in result
assert "classDef externalStyle fill:#BEDAE4" in result
# Verify it contains nodes and edges
lines = result.split("\n")
node_lines = [line for line in lines if ":::" in line]
edge_lines = [line for line in lines if " --> " in line]
assert len(node_lines) > 0
assert len(edge_lines) > 0
uv.lock
Generated
+3 -1
View File
@@ -1654,6 +1654,7 @@ dependencies = [
[package.dev-dependencies]
dev = [
{ name = "pre-commit" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-cov" },
@@ -1668,6 +1669,7 @@ requires-dist = [
[package.metadata.requires-dev]
dev = [
{ name = "pre-commit", specifier = ">=4.3.0" },
{ name = "pytest", specifier = ">=8.4.2" },
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
{ name = "pytest-cov", specifier = ">=6.1.1" },
@@ -1675,7 +1677,7 @@ dev = [
[[package]]
name = "llama-index-workflows"
version = "2.11.6"
version = "2.11.7"
source = { editable = "packages/llama-index-workflows" }
dependencies = [
{ name = "eval-type-backport", marker = "python_full_version < '3.10'" },