diff --git a/python/packages/main/agent_framework/_workflow/_viz.py b/python/packages/main/agent_framework/_workflow/_viz.py index 34ecd7e2be..37eafdd6d3 100644 --- a/python/packages/main/agent_framework/_workflow/_viz.py +++ b/python/packages/main/agent_framework/_workflow/_viz.py @@ -89,25 +89,33 @@ class WorkflowViz: 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}") + try: + 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("")) - # 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}" + except graphviz.backend.execute.ExecutableNotFound as e: + raise ImportError( + "The graphviz executables are not found. The graphviz Python package is installed, but the " + "graphviz executables (dot, neato, etc.) are not available on your system's PATH. " + "Install graphviz executables: sudo apt-get install graphviz on Debian/Ubuntu, " + "brew install graphviz on macOS, or download from https://graphviz.org/download/ for other platforms." + ) from e def save_svg(self, filename: str) -> str: """Convenience method to save as SVG. diff --git a/python/packages/main/tests/workflow/test_viz.py b/python/packages/main/tests/workflow/test_viz.py index d445343590..7c76eeeb03 100644 --- a/python/packages/main/tests/workflow/test_viz.py +++ b/python/packages/main/tests/workflow/test_viz.py @@ -186,6 +186,34 @@ def test_workflow_viz_unsupported_format(): viz.export(format="invalid") # type: ignore +def test_workflow_viz_graphviz_binary_not_found(): + """Test that missing graphviz binary raises ImportError with helpful message.""" + import unittest.mock + + # Skip test if graphviz package is not available + pytest.importorskip("graphviz") + + executor1 = MockExecutor(id="executor1") + executor2 = MockExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + viz = WorkflowViz(workflow) + + # Mock graphviz.Source.render to raise ExecutableNotFound + with unittest.mock.patch("graphviz.Source") as mock_source_class: + mock_source = unittest.mock.MagicMock() + mock_source_class.return_value = mock_source + + # Import the ExecutableNotFound exception for the test + from graphviz.backend.execute import ExecutableNotFound + + mock_source.render.side_effect = ExecutableNotFound("failed to execute PosixPath('dot')") + + # Test that the proper ImportError is raised with helpful message + with pytest.raises(ImportError, match="The graphviz executables are not found"): + viz.export(format="svg") + + def test_workflow_viz_conditional_edge(): """Test that conditional edges are rendered dashed with a label.""" start = MockExecutor(id="start")