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>
2026-06-16 21:04:09 +08:00 · 2025-08-19 17:07:49 -07:00
parent baaa9c0aee
commit 6900d3456c
11 changed files with 617 additions and 4 deletions
@@ -28,6 +28,7 @@ _IMPORTS = [
    "RequestInfoMessage",
    "WorkflowRunResult",
    "Workflow",
+    "WorkflowViz",
    "FileCheckpointStorage",
    "InMemoryCheckpointStorage",
    "CheckpointStorage",
@@ -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",
 ]
@@ -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",
@@ -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
@@ -26,6 +26,11 @@ dependencies = [
    "agent-framework",
 ]

+[project.optional-dependencies]
+viz = [
+    "graphviz>=0.20.0",
+]
+
 [tool.uv]
 prerelease = "if-necessary-or-explicit"
 environments = [
@@ -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: "<id>" [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: <id>((fan-in)) -> extract <id>
+    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
@@ -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/).
@@ -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
@@ -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
@@ -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()
@@ -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"