From 6900d3456c7568f6bba9e2471eda5012fc587bc3 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Tue, 19 Aug 2025 17:07:49 -0700 Subject: [PATCH] Python: Workflow visualization. (#372) * Workflow visualization. * fix typing * address comments * update all samples * Update fan in edge group visual * fix quality check * Update python/samples/getting_started/workflow/step_06_map_reduce.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * address comment * Catch up with changes in workflow; add to_mermaid method * Update examples * fix test * add installation guide to error messages * Remove visualization for samples except for one. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com> --- .../main/agent_framework/workflow/__init__.py | 1 + .../agent_framework/workflow/__init__.pyi | 2 + .../agent_framework_workflow/__init__.py | 2 + .../workflow/agent_framework_workflow/_viz.py | 256 ++++++++++++++++ python/packages/workflow/pyproject.toml | 5 + python/packages/workflow/tests/test_viz.py | 285 ++++++++++++++++++ .../getting_started/workflow/README.md | 11 + .../step_01_simple_workflow_sequential.py | 8 +- ...a_simple_workflow_sequential_non_stream.py | 8 +- ...> step_06_map_reduce_and_visualization.py} | 21 ++ python/uv.lock | 22 +- 11 files changed, 617 insertions(+), 4 deletions(-) create mode 100644 python/packages/workflow/agent_framework_workflow/_viz.py create mode 100644 python/packages/workflow/tests/test_viz.py create mode 100644 python/samples/getting_started/workflow/README.md rename python/samples/getting_started/workflow/{step_06_map_reduce.py => step_06_map_reduce_and_visualization.py} (93%) diff --git a/python/packages/main/agent_framework/workflow/__init__.py b/python/packages/main/agent_framework/workflow/__init__.py index 971b7fb004..89f24b0aec 100644 --- a/python/packages/main/agent_framework/workflow/__init__.py +++ b/python/packages/main/agent_framework/workflow/__init__.py @@ -28,6 +28,7 @@ _IMPORTS = [ "RequestInfoMessage", "WorkflowRunResult", "Workflow", + "WorkflowViz", "FileCheckpointStorage", "InMemoryCheckpointStorage", "CheckpointStorage", diff --git a/python/packages/main/agent_framework/workflow/__init__.pyi b/python/packages/main/agent_framework/workflow/__init__.pyi index e36c6ee461..6ee7a08a69 100644 --- a/python/packages/main/agent_framework/workflow/__init__.pyi +++ b/python/packages/main/agent_framework/workflow/__init__.pyi @@ -26,6 +26,7 @@ from agent_framework_workflow import ( WorkflowEvent, WorkflowRunResult, WorkflowStartedEvent, + WorkflowViz, __version__, handler, ) @@ -56,6 +57,7 @@ __all__ = [ "WorkflowEvent", "WorkflowRunResult", "WorkflowStartedEvent", + "WorkflowViz", "__version__", "handler", ] diff --git a/python/packages/workflow/agent_framework_workflow/__init__.py b/python/packages/workflow/agent_framework_workflow/__init__.py index 8196d5f767..7bacabee0f 100644 --- a/python/packages/workflow/agent_framework_workflow/__init__.py +++ b/python/packages/workflow/agent_framework_workflow/__init__.py @@ -45,6 +45,7 @@ from ._validation import ( WorkflowValidationError, validate_workflow_graph, ) +from ._viz import WorkflowViz from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult from ._workflow_context import WorkflowContext @@ -91,6 +92,7 @@ __all__ = [ "WorkflowRunResult", "WorkflowStartedEvent", "WorkflowValidationError", + "WorkflowViz", "__version__", "handler", "validate_workflow_graph", diff --git a/python/packages/workflow/agent_framework_workflow/_viz.py b/python/packages/workflow/agent_framework_workflow/_viz.py new file mode 100644 index 0000000000..77260ae10f --- /dev/null +++ b/python/packages/workflow/agent_framework_workflow/_viz.py @@ -0,0 +1,256 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Workflow visualization module using graphviz.""" + +import hashlib +import re +import tempfile +from pathlib import Path +from typing import Literal + +from ._edge import FanInEdgeGroup +from ._workflow import Workflow + + +class WorkflowViz: + """A class for visualizing workflows using graphviz.""" + + def __init__(self, workflow: Workflow): + """Initialize the WorkflowViz with a workflow. + + Args: + workflow: The workflow to visualize. + """ + self._workflow = workflow + + def to_digraph(self) -> str: + """Export the workflow as a DOT format digraph string. + + Returns: + A string representation of the workflow in DOT format. + """ + lines = ["digraph Workflow {"] + lines.append(" rankdir=TD;") # Top to bottom layout + lines.append(" node [shape=box, style=filled, fillcolor=lightblue];") + lines.append(" edge [color=black, arrowhead=vee];") + lines.append("") + + # Add start executor with special styling + start_executor = self._workflow.start_executor + lines.append(f' "{start_executor.id}" [fillcolor=lightgreen, label="{start_executor.id}\\n(Start)"];') + + # Add all other executors + for executor in self._workflow.executors: + if executor.id != start_executor.id: + lines.append(f' "{executor.id}" [label="{executor.id}"];') + + # Build shared structures + fan_in_nodes = self._compute_fan_in_descriptors() # (node_id, sources, target) + normal_edges = self._compute_normal_edges() # (src, tgt, is_conditional) + + if fan_in_nodes: + lines.append("") + for node_id, _, _ in fan_in_nodes: + lines.append(f' "{node_id}" [shape=ellipse, fillcolor=lightgoldenrod, label="fan-in"];') + + # Route fan-in via intermediate nodes + for node_id, sources, target in fan_in_nodes: + for src in sources: + lines.append(f' "{src}" -> "{node_id}";') + lines.append(f' "{node_id}" -> "{target}";') + + # Draw normal edges + for src, tgt, is_cond in normal_edges: + edge_attr = ' [style=dashed, label="conditional"]' if is_cond else "" + lines.append(f' "{src}" -> "{tgt}"{edge_attr};') + + lines.append("}") + return "\n".join(lines) + + def export(self, format: Literal["svg", "png", "pdf", "dot"] = "svg", filename: str | None = None) -> str: + """Export the workflow visualization to a file or return the file path. + + Args: + format: The output format. Supported formats: 'svg', 'png', 'pdf', 'dot'. + filename: Optional filename to save the output. If None, creates a temporary file. + + Returns: + The path to the saved file. + + Raises: + ImportError: If graphviz is not installed. + ValueError: If an unsupported format is specified. + """ + # Validate format first + if format not in ["svg", "png", "pdf", "dot"]: + raise ValueError(f"Unsupported format: {format}. Supported formats: svg, png, pdf, dot") + + if format == "dot": + content = self.to_digraph() + if filename: + with open(filename, "w", encoding="utf-8") as f: + f.write(content) + return filename + # Create temporary file for dot format + with tempfile.NamedTemporaryFile(mode="w", suffix=".dot", delete=False, encoding="utf-8") as temp_file: + temp_file.write(content) + return temp_file.name + + try: + import graphviz # type: ignore + except ImportError as e: + raise ImportError( + "viz extra is required for export. Install it with: pip install agent-framework-workflow[viz]. " + "You also need to install graphviz separately. E.g., sudo apt-get install graphviz on Debian/Ubuntu " + "or brew install graphviz on macOS. See https://graphviz.org/download/ for details." + ) from e + + # Create a temporary graphviz Source object + dot_content = self.to_digraph() + source = graphviz.Source(dot_content) + + if filename: + # Save to specified file + output_path = Path(filename) + if output_path.suffix and output_path.suffix[1:] != format: + raise ValueError(f"File extension {output_path.suffix} doesn't match format {format}") + + # Remove extension if present since graphviz.render() adds it + base_name = str(output_path.with_suffix("")) + source.render(base_name, format=format, cleanup=True) + + # Return the actual filename with extension + return f"{base_name}.{format}" + # Create temporary file + with tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False) as temp_file: + temp_path = Path(temp_file.name) + base_name = str(temp_path.with_suffix("")) + + source.render(base_name, format=format, cleanup=True) + return f"{base_name}.{format}" + + def save_svg(self, filename: str) -> str: + """Convenience method to save as SVG. + + Args: + filename: The filename to save the SVG file. + + Returns: + The path to the saved SVG file. + """ + return self.export(format="svg", filename=filename) + + def save_png(self, filename: str) -> str: + """Convenience method to save as PNG. + + Args: + filename: The filename to save the PNG file. + + Returns: + The path to the saved PNG file. + """ + return self.export(format="png", filename=filename) + + def save_pdf(self, filename: str) -> str: + """Convenience method to save as PDF. + + Args: + filename: The filename to save the PDF file. + + Returns: + The path to the saved PDF file. + """ + return self.export(format="pdf", filename=filename) + + def to_mermaid(self) -> str: + """Export the workflow as a Mermaid flowchart string. + + Returns: + A string representation of the workflow in Mermaid flowchart syntax. + """ + + def _san(s: str) -> str: + """Sanitize an ID for Mermaid (alphanumeric and underscore, start with letter).""" + s2 = re.sub(r"[^0-9A-Za-z_]", "_", s) + if not s2 or not s2[0].isalpha(): + s2 = f"n_{s2}" + return s2 + + lines: list[str] = ["flowchart TD"] + + # Nodes + start_executor = self._workflow.start_executor + start_id = _san(start_executor.id) + # End statements with semicolons for better compatibility and quote labels for special chars + lines.append(f' {start_id}["{start_executor.id} (Start)"];') + + for executor in self._workflow.executors: + if executor.id == start_executor.id: + continue + eid = _san(executor.id) + lines.append(f' {eid}["{executor.id}"];') + + # Build shared structures + fan_in_nodes_dot = self._compute_fan_in_descriptors() # uses DOT node ids + # Convert DOT-style node ids to Mermaid-safe ones + fan_in_nodes: list[tuple[str, list[str], str]] = [] + for dot_node_id, sources, target in fan_in_nodes_dot: + digest = dot_node_id.split("::")[-1] + fan_node_id = f"fan_in__{_san(target)}__{digest}" + fan_in_nodes.append((fan_node_id, sources, target)) + + for fan_node_id, _, _ in fan_in_nodes: + # Use double parentheses to make it circular in Mermaid + # (Keep this line without a trailing semicolon to match existing tests.) + lines.append(f" {fan_node_id}((fan-in))") + + # Fan-in edges + for fan_node_id, sources, target in fan_in_nodes: + for s in sources: + lines.append(f" {_san(s)} --> {fan_node_id};") + lines.append(f" {fan_node_id} --> {_san(target)};") + + # Normal edges + for src, tgt, is_cond in self._compute_normal_edges(): + s = _san(src) + t = _san(tgt) + if is_cond: + lines.append(f" {s} -. conditional .-> {t};") + else: + lines.append(f" {s} --> {t};") + + return "\n".join(lines) + + # region Private helpers + + def _fan_in_digest(self, target: str, sources: list[str]) -> str: + sources_sorted = sorted(sources) + return hashlib.sha256((target + "|" + "|".join(sources_sorted)).encode("utf-8")).hexdigest()[:8] + + def _compute_fan_in_descriptors(self) -> list[tuple[str, list[str], str]]: + """Return list of (node_id, sources, target) for fan-in groups. + + node_id is DOT-oriented: fan_in::target::digest + """ + result: list[tuple[str, list[str], str]] = [] + for group in self._workflow.edge_groups: + if isinstance(group, FanInEdgeGroup): + target = group.target_executors[0].id + sources = [src.id for src in group.source_executors] + digest = self._fan_in_digest(target, sources) + node_id = f"fan_in::{target}::{digest}" + result.append((node_id, sorted(sources), target)) + return result + + def _compute_normal_edges(self) -> list[tuple[str, str, bool]]: + """Return list of (source_id, target_id, is_conditional) for non-fan-in groups.""" + edges: list[tuple[str, str, bool]] = [] + for group in self._workflow.edge_groups: + if isinstance(group, FanInEdgeGroup): + continue + for edge in group.edges: + is_cond = getattr(edge, "_condition", None) is not None + edges.append((edge.source_id, edge.target_id, is_cond)) + return edges + + # endregion diff --git a/python/packages/workflow/pyproject.toml b/python/packages/workflow/pyproject.toml index d9358c80e3..5ffab1a8cb 100644 --- a/python/packages/workflow/pyproject.toml +++ b/python/packages/workflow/pyproject.toml @@ -26,6 +26,11 @@ dependencies = [ "agent-framework", ] +[project.optional-dependencies] +viz = [ + "graphviz>=0.20.0", +] + [tool.uv] prerelease = "if-necessary-or-explicit" environments = [ diff --git a/python/packages/workflow/tests/test_viz.py b/python/packages/workflow/tests/test_viz.py new file mode 100644 index 0000000000..72c65caa69 --- /dev/null +++ b/python/packages/workflow/tests/test_viz.py @@ -0,0 +1,285 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for the workflow visualization module.""" + +import pytest +from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowContext, WorkflowViz, handler + + +class MockExecutor(Executor): + """A mock executor for testing purposes.""" + + @handler(output_types=[str]) + async def mock_handler(self, message: str, ctx: WorkflowContext) -> None: + """A mock handler that does nothing.""" + pass + + +class ListStrTargetExecutor(Executor): + """A mock executor that accepts a list of strings (for fan-in targets).""" + + @handler + async def handle(self, message: list[str], ctx: WorkflowContext) -> None: # type: ignore[type-arg] + pass + + +def test_workflow_viz_to_digraph(): + """Test that WorkflowViz can generate a DOT digraph.""" + # Create a simple workflow + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + viz = WorkflowViz(workflow) + dot_content = viz.to_digraph() + + # Check that the DOT content contains expected elements + assert "digraph Workflow {" in dot_content + assert '"executor1"' in dot_content + assert '"executor2"' in dot_content + assert '"executor1" -> "executor2"' in dot_content + assert "fillcolor=lightgreen" in dot_content # Start executor styling + assert "(Start)" in dot_content + + +def test_workflow_viz_export_dot(): + """Test exporting workflow as DOT format.""" + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + viz = WorkflowViz(workflow) + + # Test export without filename (returns temporary file path) + file_path = viz.export(format="dot") + assert file_path.endswith(".dot") + + with open(file_path, encoding="utf-8") as f: + content = f.read() + + assert "digraph Workflow {" in content + assert '"executor1" -> "executor2"' in content + + +def test_workflow_viz_export_dot_with_filename(tmp_path): + """Test exporting workflow as DOT format with specified filename.""" + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + viz = WorkflowViz(workflow) + + # Test export with filename + output_file = tmp_path / "test_workflow.dot" + result_path = viz.export(format="dot", filename=str(output_file)) + + assert result_path == str(output_file) + assert output_file.exists() + + content = output_file.read_text(encoding="utf-8") + assert "digraph Workflow {" in content + assert '"executor1" -> "executor2"' in content + + +def test_workflow_viz_complex_workflow(): + """Test visualization of a more complex workflow.""" + executor1 = MockExecutor(id="start") + executor2 = MockExecutor(id="middle1") + executor3 = MockExecutor(id="middle2") + executor4 = MockExecutor(id="end") + + workflow = ( + WorkflowBuilder() + .add_edge(executor1, executor2) + .add_edge(executor1, executor3) + .add_edge(executor2, executor4) + .add_edge(executor3, executor4) + .set_start_executor(executor1) + .build() + ) + + viz = WorkflowViz(workflow) + dot_content = viz.to_digraph() + + # Check all executors are present + assert '"start"' in dot_content + assert '"middle1"' in dot_content + assert '"middle2"' in dot_content + assert '"end"' in dot_content + + # Check all edges are present + assert '"start" -> "middle1"' in dot_content + assert '"start" -> "middle2"' in dot_content + assert '"middle1" -> "end"' in dot_content + assert '"middle2" -> "end"' in dot_content + + # Check start executor has special styling + assert "fillcolor=lightgreen" in dot_content + + +@pytest.mark.skipif(True, reason="Requires graphviz to be installed") +def test_workflow_viz_export_svg(): + """Test exporting workflow as SVG format. Skipped unless graphviz is available.""" + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + viz = WorkflowViz(workflow) + + try: + file_path = viz.export(format="svg") + assert file_path.endswith(".svg") + except ImportError: + pytest.skip("graphviz not available") + + +def test_workflow_viz_unsupported_format(): + """Test that unsupported formats raise ValueError.""" + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + viz = WorkflowViz(workflow) + + with pytest.raises(ValueError, match="Unsupported format: invalid"): + viz.export(format="invalid") # type: ignore + + +def test_workflow_viz_conditional_edge(): + """Test that conditional edges are rendered dashed with a label.""" + start = MockExecutor(id="start") + mid = MockExecutor(id="mid") + end = MockExecutor(id="end") + + # Condition that is never used during viz, but presence should mark the edge + def only_if_foo(msg: str) -> bool: # pragma: no cover - simple predicate + return msg == "foo" + + wf = ( + WorkflowBuilder() + .add_edge(start, mid, condition=only_if_foo) + .add_edge(mid, end) + .set_start_executor(start) + .build() + ) + + dot = WorkflowViz(wf).to_digraph() + + # Conditional edge should be dashed and labeled + assert '"start" -> "mid" [style=dashed, label="conditional"];' in dot + # Non-conditional edge should be plain + assert '"mid" -> "end"' in dot + assert '"mid" -> "end" [style=dashed' not in dot + + +def test_workflow_viz_fan_in_edge_group(): + """Test that fan-in edges render an intermediate node with label and routed edges.""" + start = MockExecutor(id="start") + s1 = MockExecutor(id="s1") + s2 = MockExecutor(id="s2") + t = ListStrTargetExecutor(id="t") + + # Build a connected workflow: start fans out to s1 and s2, which then fan-in to t + wf = ( + WorkflowBuilder() + .add_fan_out_edges(start, [s1, s2]) + .add_fan_in_edges([s1, s2], t) + .set_start_executor(start) + .build() + ) + + dot = WorkflowViz(wf).to_digraph() + + # There should be a single fan-in node with special styling and label + lines = [line.strip() for line in dot.splitlines()] + fan_in_lines = [line for line in lines if "shape=ellipse" in line and 'label="fan-in"' in line] + assert len(fan_in_lines) == 1 + + # Extract the intermediate node id from the line: "" [shape=ellipse, ... label="fan-in"]; + fan_in_line = fan_in_lines[0] + first_quote = fan_in_line.find('"') + second_quote = fan_in_line.find('"', first_quote + 1) + assert first_quote != -1 and second_quote != -1 + fan_in_node_id = fan_in_line[first_quote + 1 : second_quote] + assert fan_in_node_id # non-empty + + # Edges should be routed through the intermediate node, not direct to target + assert f'"s1" -> "{fan_in_node_id}";' in dot + assert f'"s2" -> "{fan_in_node_id}";' in dot + assert f'"{fan_in_node_id}" -> "t";' in dot + + # Ensure direct edges are not present + assert '"s1" -> "t"' not in dot + assert '"s2" -> "t"' not in dot + + +def test_workflow_viz_to_mermaid_basic(): + """Mermaid: basic workflow nodes and edge are present with start label.""" + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + mermaid = WorkflowViz(workflow).to_mermaid() + + # Start node and normal node + assert 'executor1["executor1 (Start)"]' in mermaid + assert 'executor2["executor2"]' in mermaid + # Edge uses sanitized ids (same as ids here) + assert "executor1 --> executor2" in mermaid + + +def test_workflow_viz_mermaid_conditional_edge(): + """Mermaid: conditional edges are dotted with a label.""" + start = MockExecutor(id="start") + mid = MockExecutor(id="mid") + + def only_if_foo(msg: str) -> bool: # pragma: no cover - simple predicate + return msg == "foo" + + wf = WorkflowBuilder().add_edge(start, mid, condition=only_if_foo).set_start_executor(start).build() + mermaid = WorkflowViz(wf).to_mermaid() + + assert "start -. conditional .-> mid" in mermaid + + +def test_workflow_viz_mermaid_fan_in_edge_group(): + """Mermaid: fan-in uses an intermediate node and routes edges via it.""" + start = MockExecutor(id="start") + s1 = MockExecutor(id="s1") + s2 = MockExecutor(id="s2") + t = ListStrTargetExecutor(id="t") + + wf = ( + WorkflowBuilder() + .add_fan_out_edges(start, [s1, s2]) + .add_fan_in_edges([s1, s2], t) + .set_start_executor(start) + .build() + ) + + mermaid = WorkflowViz(wf).to_mermaid() + lines = [line.strip() for line in mermaid.splitlines()] + # Find the fan-in node (line ends with ((fan-in))) + fan_lines = [ln for ln in lines if ln.endswith("((fan-in))")] + assert len(fan_lines) == 1 + fan_line = fan_lines[0] + # fan_in node is emitted as: ((fan-in)) -> extract + token = fan_line.strip() + suffix = "((fan-in))" + assert token.endswith(suffix) + fan_node_id = token[: -len(suffix)] + assert fan_node_id + + # Ensure routing via the intermediate node + assert f"s1 --> {fan_node_id}" in mermaid + assert f"s2 --> {fan_node_id}" in mermaid + assert f"{fan_node_id} --> t" in mermaid + + # Ensure direct edges to target are not present + assert "s1 --> t" not in mermaid + assert "s2 --> t" not in mermaid diff --git a/python/samples/getting_started/workflow/README.md b/python/samples/getting_started/workflow/README.md new file mode 100644 index 0000000000..f35113203e --- /dev/null +++ b/python/samples/getting_started/workflow/README.md @@ -0,0 +1,11 @@ +# Workflow Getting Started Samples + +## Installation + +You can install the workflow package with visualization dependency: + +```bash +pip install agent-framework-workflow[viz] +``` + +To export visualization images you also need to [install GraphViz](https://graphviz.org/download/). diff --git a/python/samples/getting_started/workflow/step_01_simple_workflow_sequential.py b/python/samples/getting_started/workflow/step_01_simple_workflow_sequential.py index ff887598d3..492676ebdc 100644 --- a/python/samples/getting_started/workflow/step_01_simple_workflow_sequential.py +++ b/python/samples/getting_started/workflow/step_01_simple_workflow_sequential.py @@ -2,7 +2,13 @@ import asyncio -from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler +from agent_framework.workflow import ( + Executor, + WorkflowBuilder, + WorkflowCompletedEvent, + WorkflowContext, + handler, +) """ The following sample demonstrates a basic workflow with two executors diff --git a/python/samples/getting_started/workflow/step_01a_simple_workflow_sequential_non_stream.py b/python/samples/getting_started/workflow/step_01a_simple_workflow_sequential_non_stream.py index d38569fcfc..37662c1d0d 100644 --- a/python/samples/getting_started/workflow/step_01a_simple_workflow_sequential_non_stream.py +++ b/python/samples/getting_started/workflow/step_01a_simple_workflow_sequential_non_stream.py @@ -2,7 +2,13 @@ import asyncio -from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler +from agent_framework.workflow import ( + Executor, + WorkflowBuilder, + WorkflowCompletedEvent, + WorkflowContext, + handler, +) """ The following sample demonstrates a basic workflow with two executors diff --git a/python/samples/getting_started/workflow/step_06_map_reduce.py b/python/samples/getting_started/workflow/step_06_map_reduce_and_visualization.py similarity index 93% rename from python/samples/getting_started/workflow/step_06_map_reduce.py rename to python/samples/getting_started/workflow/step_06_map_reduce_and_visualization.py index 929cb5f4c7..aad2732337 100644 --- a/python/samples/getting_started/workflow/step_06_map_reduce.py +++ b/python/samples/getting_started/workflow/step_06_map_reduce_and_visualization.py @@ -12,6 +12,7 @@ from agent_framework.workflow import ( WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, + WorkflowViz, handler, ) @@ -23,6 +24,8 @@ to a final count per word. Intermediate results are stored in a temporary directory, and the final results are written to a file in the same directory. + +This sample also shows how you can visualize a workflow using `WorkflowViz`. """ # Define the temporary directory for storing intermediate results @@ -286,6 +289,24 @@ async def main(): .build() ) + # Step 2.5: Visualize the workflow (optional) + print("🎨 Generating workflow visualization...") + viz = WorkflowViz(workflow) + # Print out the mermaid string. + print("🧜 Mermaid string: \n=======") + print(viz.to_mermaid()) + print("=======") + # Print out the DiGraph string. + print("📊 DiGraph string: \n=======") + print(viz.to_digraph()) + print("=======") + try: + # Export the DiGraph visualization as SVG. + svg_file = viz.export(format="svg") + print(f"🖼️ SVG file saved to: {svg_file}") + except ImportError: + print("💡 Tip: Install 'viz' extra to export workflow visualization: pip install agent-framework-workflow[viz]") + # Step 3: Open the text file and read its content. async with aiofiles.open(os.path.join(DIR, "resources", "long_text.txt"), "r") as f: raw_text = await f.read() diff --git a/python/uv.lock b/python/uv.lock index 4a407af285..e9827015c8 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -204,8 +204,17 @@ dependencies = [ { name = "agent-framework", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] +[package.optional-dependencies] +viz = [ + { name = "graphviz", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] + [package.metadata] -requires-dist = [{ name = "agent-framework", editable = "packages/main" }] +requires-dist = [ + { name = "agent-framework", editable = "packages/main" }, + { name = "graphviz", marker = "extra == 'viz'", specifier = ">=0.20.0" }, +] +provides-extras = ["viz"] [[package]] name = "aiofiles" @@ -894,7 +903,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -1031,6 +1040,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/45/b82e3c16be2182bff01179db177fe144d58b5dc787a7d4492c6ed8b9317f/frozenlist-1.7.0-py3-none-any.whl", hash = "sha256:9a5af342e34f7e97caf8c995864c7a396418ae2859cc6fdf1b1073020d516a7e", size = 13106, upload-time = "2025-06-09T23:02:34.204Z" }, ] +[[package]] +name = "graphviz" +version = "0.21" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f8/b3/3ac91e9be6b761a4b30d66ff165e54439dcd48b83f4e20d644867215f6ca/graphviz-0.21.tar.gz", hash = "sha256:20743e7183be82aaaa8ad6c93f8893c923bd6658a04c32ee115edb3c8a835f78", size = 200434, upload-time = "2025-06-15T09:35:05.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/4c/e0ce1ef95d4000ebc1c11801f9b944fa5910ecc15b5e351865763d8657f8/graphviz-0.21-py3-none-any.whl", hash = "sha256:54f33de9f4f911d7e84e4191749cac8cc5653f815b06738c54db9a15ab8b1e42", size = 47300, upload-time = "2025-06-15T09:35:04.433Z" }, +] + [[package]] name = "greenlet" version = "3.2.4"