diff --git a/src/crewai/cli/plot_flow.py b/src/crewai/cli/plot_flow.py index 848c55d69..b8afae4c6 100644 --- a/src/crewai/cli/plot_flow.py +++ b/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) diff --git a/src/crewai/crew.py b/src/crewai/crew.py index f1c33f637..3347823ca 100644 --- a/src/crewai/crew.py +++ b/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. diff --git a/tests/cli/test_plot_flow.py b/tests/cli/test_plot_flow.py new file mode 100644 index 000000000..0a4ed0913 --- /dev/null +++ b/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() diff --git a/tests/test_crew_plot.py b/tests/test_crew_plot.py new file mode 100644 index 000000000..dc833a015 --- /dev/null +++ b/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)