Fix plot functionality (#2638)

Co-Authored-By: Joe Moura <joao@crewai.com>

src/crewai/cli/plot_flow.py

@@ -1,23 +1,29 @@
-import subprocess
+import os
+import importlib.util
+import sys
 
 import click
 
 
 def plot_flow() -> None:
     """
-    Plot the flow by running a command in the UV environment.
+    Plot the flow by finding and importing the plot function from the project's main.py file.
     """
-    command = ["uv", "run", "plot"]
-
+    main_file_path = os.path.join(os.getcwd(), "main.py")
+    
     try:
-        result = subprocess.run(command, capture_output=False, text=True, check=True)
-
-        if result.stderr:
-            click.echo(result.stderr, err=True)
-
-    except subprocess.CalledProcessError as e:
-        click.echo(f"An error occurred while plotting the flow: {e}", err=True)
-        click.echo(e.output, err=True)
-
+        if os.path.exists(main_file_path):
+            spec = importlib.util.spec_from_file_location("main", main_file_path)
+            main = importlib.util.module_from_spec(spec)
+            sys.modules["main"] = main
+            spec.loader.exec_module(main)
+            
+            if hasattr(main, "plot"):
+                main.plot()
+            else:
+                click.echo("Error: No plot function found in main.py", err=True)
+        else:
+            click.echo("Error: Could not find main.py in the current directory", err=True)
+            
     except Exception as e:
-        click.echo(f"An unexpected error occurred: {e}", err=True)
+        click.echo(f"An error occurred while plotting the flow: {e}", err=True)

src/crewai/crew.py

@@ -1328,6 +1328,19 @@ class Crew(BaseModel):
 
     def __repr__(self):
         return f"Crew(id={self.id}, process={self.process}, number_of_agents={len(self.agents)}, number_of_tasks={len(self.tasks)})"
+        
+    def plot(self, filename: str = "crewai_crew") -> None:
+        """
+        This method is not implemented for Crew objects.
+        Plot functionality is available only for Flow objects.
+        
+        To visualize a Flow, use flow_instance.plot() instead.
+        """
+        raise NotImplementedError(
+            "The plot method is not available for Crew objects. "
+            "Plot functionality is only available for Flow objects. "
+            "Please use flow_instance.plot() instead."
+        )
 
     def reset_memories(self, command_type: str) -> None:
         """Reset specific or all memories for the crew.

tests/cli/test_plot_flow.py

@@ -0,0 +1,37 @@
+import os
+import tempfile
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from crewai.cli.plot_flow import plot_flow
+
+
+class TestPlotFlow:
+    def test_plot_flow_no_main_file(self):
+        """Test plot_flow when main.py doesn't exist."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            original_dir = os.getcwd()
+            try:
+                os.chdir(temp_dir)
+                with patch("click.echo") as mock_echo:
+                    plot_flow()
+                    mock_echo.assert_called_with(
+                        "Error: Could not find main.py in the current directory", err=True
+                    )
+            finally:
+                os.chdir(original_dir)
+
+    def test_plot_flow_with_main_file(self):
+        """Test plot_flow with a mock main.py that has plot function."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            with patch("importlib.util.spec_from_file_location") as mock_spec_from_file:
+                with patch("click.echo"):
+                    mock_module = MagicMock()
+                    mock_spec = MagicMock()
+                    mock_spec_from_file.return_value = mock_spec
+                    mock_spec.loader = MagicMock()
+                    
+                    with patch("os.path.exists", return_value=True):
+                        plot_flow()
+                        mock_spec.loader.exec_module.assert_called_once()

tests/test_crew_plot.py

@@ -0,0 +1,20 @@
+import pytest
+
+from crewai.crew import Crew
+from crewai.agent import Agent
+
+
+def test_crew_plot_method_not_implemented():
+    """Test that trying to plot a Crew raises the correct error."""
+    agent = Agent(
+        role="Test Agent",
+        goal="Test Goal",
+        backstory="Test Backstory"
+    )
+    crew = Crew(agents=[agent], tasks=[])
+    
+    with pytest.raises(NotImplementedError) as excinfo:
+        crew.plot()
+    
+    assert "plot method is not available for Crew objects" in str(excinfo.value)
+    assert "Plot functionality is only available for Flow objects" in str(excinfo.value)