diff --git a/src/crewai/cli/cli.py b/src/crewai/cli/cli.py index 22e1aad3c..de6160ba6 100644 --- a/src/crewai/cli/cli.py +++ b/src/crewai/cli/cli.py @@ -12,12 +12,14 @@ from crewai.memory.storage.kickoff_task_outputs_storage import ( from .authentication.main import AuthenticationCommand from .deploy.main import DeployCommand -from .tools.main import ToolCommand from .evaluate_crew import evaluate_crew from .install_crew import install_crew +from .plot_flow import plot_flow from .replay_from_task import replay_task_command from .reset_memories_command import reset_memories_command from .run_crew import run_crew +from .run_flow import run_flow +from .tools.main import ToolCommand from .train_crew import train_crew @@ -273,5 +275,25 @@ def tool_publish(is_public: bool): tool_cmd.publish(is_public) +@crewai.group() +def flow(): + """Flow related commands.""" + pass + + +@flow.command(name="run") +def flow_run(): + """Run the Flow.""" + click.echo("Running the Flow") + run_flow() + + +@flow.command(name="plot") +def flow_plot(): + """Plot the Flow.""" + click.echo("Plotting the Flow") + plot_flow() + + if __name__ == "__main__": crewai() diff --git a/src/crewai/cli/plot_flow.py b/src/crewai/cli/plot_flow.py new file mode 100644 index 000000000..bb7b2052f --- /dev/null +++ b/src/crewai/cli/plot_flow.py @@ -0,0 +1,23 @@ +import subprocess + +import click + + +def plot_flow() -> None: + """ + Plot the flow by running a command in the Poetry environment. + """ + command = ["poetry", "run", "plot_flow"] + + 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) + + except Exception as e: + click.echo(f"An unexpected error occurred: {e}", err=True) diff --git a/src/crewai/cli/run_flow.py b/src/crewai/cli/run_flow.py new file mode 100644 index 000000000..3a9e72817 --- /dev/null +++ b/src/crewai/cli/run_flow.py @@ -0,0 +1,23 @@ +import subprocess + +import click + + +def run_flow() -> None: + """ + Run the flow by running a command in the Poetry environment. + """ + command = ["poetry", "run", "run_flow"] + + 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 running the flow: {e}", err=True) + click.echo(e.output, err=True) + + except Exception as e: + click.echo(f"An unexpected error occurred: {e}", err=True) diff --git a/src/crewai/cli/templates/flow/main.py b/src/crewai/cli/templates/flow/main.py index bda89065d..38d2d8736 100644 --- a/src/crewai/cli/templates/flow/main.py +++ b/src/crewai/cli/templates/flow/main.py @@ -22,8 +22,7 @@ class PoemFlow(Flow[PoemState]): def generate_poem(self): print("Generating poem") print(f"State before poem: {self.state}") - poem_crew = PoemCrew().crew() - result = poem_crew.kickoff(inputs={"sentence_count": self.state.sentence_count}) + result = PoemCrew().crew().kickoff(inputs={"sentence_count": self.state.sentence_count}) print("Poem generated", result.raw) self.state.poem = result.raw @@ -38,16 +37,28 @@ class PoemFlow(Flow[PoemState]): f.write(self.state.poem) print(f"State after save_poem: {self.state}") -async def run(): +async def run_flow(): """ Run the flow. """ poem_flow = PoemFlow() await poem_flow.kickoff() +async def plot_flow(): + """ + Plot the flow. + """ + poem_flow = PoemFlow() + poem_flow.plot() + def main(): - asyncio.run(run()) + asyncio.run(run_flow()) + + +def plot(): + asyncio.run(plot_flow()) + if __name__ == "__main__": diff --git a/src/crewai/cli/templates/flow/pyproject.toml b/src/crewai/cli/templates/flow/pyproject.toml index 0e7083c71..4e4fd4aba 100644 --- a/src/crewai/cli/templates/flow/pyproject.toml +++ b/src/crewai/cli/templates/flow/pyproject.toml @@ -11,7 +11,8 @@ asyncio = "*" [tool.poetry.scripts] {{folder_name}} = "{{folder_name}}.main:main" -run_crew = "{{folder_name}}.main:main" +run_flow = "{{folder_name}}.main:main" +plot_flow = "{{folder_name}}.main:plot" [build-system] requires = ["poetry-core"] diff --git a/src/crewai/flow/flow.py b/src/crewai/flow/flow.py index 4dd761137..a29c9c3f9 100644 --- a/src/crewai/flow/flow.py +++ b/src/crewai/flow/flow.py @@ -6,7 +6,7 @@ from typing import Any, Callable, Dict, Generic, List, Set, Type, TypeVar, Union from pydantic import BaseModel -from crewai.flow.flow_visualizer import visualize_flow +from crewai.flow.flow_visualizer import plot_flow T = TypeVar("T", bound=Union[BaseModel, Dict[str, Any]]) @@ -268,5 +268,5 @@ class Flow(Generic[T], metaclass=FlowMeta): traceback.print_exc() - def visualize(self, filename: str = "crewai_flow_graph"): - visualize_flow(self, filename) + def plot(self, filename: str = "crewai_flow_graph"): + plot_flow(self, filename) diff --git a/src/crewai/flow/flow_visualizer.py b/src/crewai/flow/flow_visualizer.py index 468b390d1..822f192b0 100644 --- a/src/crewai/flow/flow_visualizer.py +++ b/src/crewai/flow/flow_visualizer.py @@ -15,13 +15,13 @@ from crewai.flow.visualization_utils import ( ) -class FlowVisualizer: +class FlowPlot: def __init__(self, flow): self.flow = flow self.colors = COLORS self.node_styles = NODE_STYLES - def visualize(self, filename): + def plot(self, filename): net = Network( directed=True, height="750px", @@ -61,6 +61,8 @@ class FlowVisualizer: f.write(final_html_content) print(f"Graph saved as {filename}.html") + self._cleanup_pyvis_lib() + def _generate_final_html(self, network_html): # Extract just the body content from the generated HTML current_dir = os.path.dirname(__file__) @@ -80,7 +82,18 @@ class FlowVisualizer: ) return final_html_content + def _cleanup_pyvis_lib(self): + # Clean up the generated lib folder + lib_folder = os.path.join(os.getcwd(), "lib") + try: + if os.path.exists(lib_folder) and os.path.isdir(lib_folder): + import shutil -def visualize_flow(flow, filename="flow_graph"): - visualizer = FlowVisualizer(flow) - visualizer.visualize(filename) + shutil.rmtree(lib_folder) + except Exception as e: + print(f"Error cleaning up {lib_folder}: {e}") + + +def plot_flow(flow, filename="flow_graph"): + visualizer = FlowPlot(flow) + visualizer.plot(filename)