update CLI and templates

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()

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)

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)

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__":

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"]

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)

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)