Add rollback functionality to setup wizard

HYPRGK · HYPRGK · commit 3fd301c2046a · 2024-12-01T17:34:54.000Z
Fixes AgentOps-AI#9 Add error handling and rollback functionality to the setup wizard. * Add `rollback_actions` function in `agentstack/cli/cli.py` to undo actions like directory creation and file writing. * Update `init_project_builder` in `agentstack/cli/cli.py` to catch exceptions and call `rollback_actions`. * Track created files and directories in `insert_template` for potential rollback in `agentstack/cli/cli.py`. * Add `rollback_actions` function in `agentstack/generation/agent_generation.py` and `agentstack/generation/task_generation.py` to undo actions. * Update `generate_agent` in `agentstack/generation/agent_generation.py` and `generate_task` in `agentstack/generation/task_generation.py` to catch exceptions and call `rollback_actions`. * Add tests in `tests/test_cli_loads.py` to verify rollback functionality. --- For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/AgentOps-AI/AgentStack/issues/9?shareId=XXXX-XXXX-XXXX-XXXX).
diff --git a/agentstack/cli/cli.py b/agentstack/cli/cli.py
@@ -20,86 +20,100 @@
 from .. import generation
 from ..utils import open_json_file, term_color, is_snake_case
 
+created_files = []
+created_dirs = []
 
-def init_project_builder(slug_name: Optional[str] = None, template: Optional[str] = None, use_wizard: bool = False):
-    if slug_name and not is_snake_case(slug_name):
-        print(term_color("Project name must be snake case", 'red'))
-        return
+def rollback_actions():
+    for file in created_files:
+        if os.path.exists(file):
+            os.remove(file)
+    for dir in created_dirs:
+        if os.path.exists(dir):
+            shutil.rmtree(dir)
 
-    if template is not None and use_wizard:
-        print(term_color("Template and wizard flags cannot be used together", 'red'))
-        return
-
-    template_data = None
-    if template is not None:
-        url_start = "https://"
-        if template[:len(url_start)] == url_start:
-            # template is a url
-            response = requests.get(template)
-            if response.status_code == 200:
-                template_data = response.json()
-            else:
-                print(term_color(f"Failed to fetch template data from {template}. Status code: {response.status_code}", 'red'))
-                sys.exit(1)
-        else:
-            with importlib.resources.path('agentstack.templates.proj_templates', f'{template}.json') as template_path:
-                if template_path is None:
-                    print(term_color(f"No such template {template} found", 'red'))
+def init_project_builder(slug_name: Optional[str] = None, template: Optional[str] = None, use_wizard: bool = False):
+    try:
+        if slug_name and not is_snake_case(slug_name):
+            print(term_color("Project name must be snake case", 'red'))
+            return
+
+        if template is not None and use_wizard:
+            print(term_color("Template and wizard flags cannot be used together", 'red'))
+            return
+
+        template_data = None
+        if template is not None:
+            url_start = "https://"
+            if template[:len(url_start)] == url_start:
+                # template is a url
+                response = requests.get(template)
+                if response.status_code == 200:
+                    template_data = response.json()
+                else:
+                    print(term_color(f"Failed to fetch template data from {template}. Status code: {response.status_code}", 'red'))
                     sys.exit(1)
-                template_data = open_json_file(template_path)
-
-    if template_data:
-        project_details = {
-            "name": slug_name or template_data['name'],
-            "version": "0.0.1",
-            "description": template_data['description'],
-            "author": "Name <Email>",
-            "license": "MIT"
-        }
-        framework = template_data['framework']
-        design = {
-            'agents': template_data['agents'],
-            'tasks': template_data['tasks']
-        }
-
-        tools = template_data['tools']
-
-    elif use_wizard:
-        welcome_message()
-        project_details = ask_project_details(slug_name)
-        welcome_message()
-        framework = ask_framework()
-        design = ask_design()
-        tools = ask_tools()
-
-    else:
-        welcome_message()
-        project_details = {
-            "name": slug_name or "agentstack_project",
-            "version": "0.0.1",
-            "description": "New agentstack project",
-            "author": "Name <Email>",
-            "license": "MIT"
-        }
-
-        framework = "CrewAI"  # TODO: if --no-wizard, require a framework flag
-
-        design = {
-            'agents': [],
-            'tasks': []
-        }
-
-        tools = []
-
-    log.debug(
-        f"project_details: {project_details}"
-        f"framework: {framework}"
-        f"design: {design}"
-    )
-    insert_template(project_details, framework, design, template_data)
-    for tool_data in tools:
-        generation.add_tool(tool_data['name'], agents=tool_data['agents'], path=project_details['name'])
+            else:
+                with importlib.resources.path('agentstack.templates.proj_templates', f'{template}.json') as template_path:
+                    if template_path is None:
+                        print(term_color(f"No such template {template} found", 'red'))
+                        sys.exit(1)
+                    template_data = open_json_file(template_path)
+
+        if template_data:
+            project_details = {
+                "name": slug_name or template_data['name'],
+                "version": "0.0.1",
+                "description": template_data['description'],
+                "author": "Name <Email>",
+                "license": "MIT"
+            }
+            framework = template_data['framework']
+            design = {
+                'agents': template_data['agents'],
+                'tasks': template_data['tasks']
+            }
+
+            tools = template_data['tools']
+
+        elif use_wizard:
+            welcome_message()
+            project_details = ask_project_details(slug_name)
+            welcome_message()
+            framework = ask_framework()
+            design = ask_design()
+            tools = ask_tools()
 
+        else:
+            welcome_message()
+            project_details = {
+                "name": slug_name or "agentstack_project",
+                "version": "0.0.1",
+                "description": "New agentstack project",
+                "author": "Name <Email>",
+                "license": "MIT"
+            }
+
+            framework = "CrewAI"  # TODO: if --no-wizard, require a framework flag
+
+            design = {
+                'agents': [],
+                'tasks': []
+            }
+
+            tools = []
+
+        log.debug(
+            f"project_details: {project_details}"
+            f"framework: {framework}"
+            f"design: {design}"
+        )
+        insert_template(project_details, framework, design, template_data)
+        for tool_data in tools:
+            generation.add_tool(tool_data['name'], agents=tool_data['agents'], path=project_details['name'])
+    except Exception as e:
+        print(term_color(f"An error occurred: {e}", 'red'))
+        rollback_actions()
+        sys.exit(1)
 
 def welcome_message():
     os.system("cls" if os.name == "nt" else "clear")
@@ -323,17 +337,20 @@ def insert_template(project_details: dict, framework_name: str, design: dict, te
     template_path = get_package_path() / f'templates/{framework.name}'
     with open(f"{template_path}/cookiecutter.json", "w") as json_file:
         json.dump(cookiecutter_data.to_dict(), json_file)
+        created_files.append(f"{template_path}/cookiecutter.json")
 
     # copy .env.example to .env
     shutil.copy(
         f'{template_path}/{"{{cookiecutter.project_metadata.project_slug}}"}/.env.example',
         f'{template_path}/{"{{cookiecutter.project_metadata.project_slug}}"}/.env')
+    created_files.append(f'{template_path}/{"{{cookiecutter.project_metadata.project_slug}}"}/.env')
 
     if os.path.isdir(project_details['name']):
         print(term_color(f"Directory {template_path} already exists. Please check this and try again", "red"))
         return
 
     cookiecutter(str(template_path), no_input=True, extra_context=None)
+    created_dirs.append(project_details['name'])
 
     # TODO: inits a git repo in the directory the command was run in
     # TODO: not where the project is generated. Fix this
@@ -378,4 +395,4 @@ def list_tools():
             print(f": {tool.url if tool.url else 'AgentStack default tool'}")
 
     print("\n\n✨ Add a tool with: agentstack tools add <tool_name>")
-    print("   https://docs.agentstack.sh/tools/core")
+    print("   https://docs.agentstack.sh/tools/core")
diff --git a/agentstack/generation/agent_generation.py b/agentstack/generation/agent_generation.py
@@ -6,6 +6,16 @@
 from ruamel.yaml import YAML
 from ruamel.yaml.scalarstring import FoldedScalarString
 
+created_files = []
+created_dirs = []
+
+def rollback_actions():
+    for file in created_files:
+        if os.path.exists(file):
+            os.remove(file)
+    for dir in created_dirs:
+        if os.path.exists(dir):
+            shutil.rmtree(dir)
 
 def generate_agent(
         name,
@@ -27,18 +37,19 @@ def generate_agent(
 
     framework = get_framework()
 
-    if framework == 'crewai':
-        generate_crew_agent(name, role, goal, backstory, llm)
-        print("    > Added to src/config/agents.yaml")
-    else:
-        print(f"This function is not yet implemented for {framework}")
-        return
-
-    print(f"Added agent \"{name}\" to your AgentStack project successfully!")
-
-
-
+    try:
+        if framework == 'crewai':
+            generate_crew_agent(name, role, goal, backstory, llm)
+            print("    > Added to src/config/agents.yaml")
+        else:
+            print(f"This function is not yet implemented for {framework}")
+            return
 
+        print(f"Added agent \"{name}\" to your AgentStack project successfully!")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        rollback_actions()
+        sys.exit(1)
 
 def generate_crew_agent(
         name,
@@ -83,6 +94,7 @@ def generate_crew_agent(
     # Write back to the file without altering existing content
     with open(config_path, 'w') as file:
         yaml.dump(data, file)
+        created_files.append(config_path)
 
     # Now lets add the agent to crew.py
     file_path = 'src/crew.py'
@@ -103,4 +115,4 @@ def generate_crew_agent(
 
 def get_agent_names(framework: str = 'crewai', path: str = '') -> List[str]:
     """Get only agent names from the crew file"""
-    return get_crew_components(framework, CrewComponent.AGENT, path)['agents']
+    return get_crew_components(framework, CrewComponent.AGENT, path)['agents']
diff --git a/agentstack/generation/task_generation.py b/agentstack/generation/task_generation.py
@@ -6,6 +6,16 @@
 from ruamel.yaml import YAML
 from ruamel.yaml.scalarstring import FoldedScalarString
 
+created_files = []
+created_dirs = []
+
+def rollback_actions():
+    for file in created_files:
+        if os.path.exists(file):
+            os.remove(file)
+    for dir in created_dirs:
+        if os.path.exists(dir):
+            shutil.rmtree(dir)
 
 def generate_task(
         name,
@@ -24,15 +34,19 @@ def generate_task(
 
     framework = get_framework()
 
-    if framework == 'crewai':
-        generate_crew_task(name, description, expected_output, agent)
-        print("    > Added to src/config/tasks.yaml")
-    else:
-        print(f"This function is not yet implemented for {framework}")
-        return
-
-    print(f"Added task \"{name}\" to your AgentStack project successfully!")
+    try:
+        if framework == 'crewai':
+            generate_crew_task(name, description, expected_output, agent)
+            print("    > Added to src/config/tasks.yaml")
+        else:
+            print(f"This function is not yet implemented for {framework}")
+            return
 
+        print(f"Added task \"{name}\" to your AgentStack project successfully!")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        rollback_actions()
+        sys.exit(1)
 
 def generate_crew_task(
         name,
@@ -74,6 +88,7 @@ def generate_crew_task(
     # Write back to the file without altering existing content
     with open(config_path, 'w') as file:
         yaml.dump(data, file)
+        created_files.append(config_path)
 
     # Add task to crew.py
     file_path = 'src/crew.py'
@@ -91,4 +106,4 @@ def generate_crew_task(
 
 def get_task_names(framework: str, path: str = '') -> List[str]:
     """Get only task names from the crew file"""
-    return get_crew_components(framework, CrewComponent.TASK, path)['tasks']
+    return get_crew_components(framework, CrewComponent.TASK, path)['tasks']
diff --git a/tests/test_cli_loads.py b/tests/test_cli_loads.py
@@ -3,6 +3,7 @@
 import unittest
 from pathlib import Path
 import shutil
+import os
 
 
 class TestAgentStackCLI(unittest.TestCase):
@@ -44,6 +45,25 @@ def test_init_command(self):
         # Clean up
         shutil.rmtree(test_dir)
 
+    def test_rollback_on_error(self):
+        """Test rollback functionality when an error occurs during project initialization."""
+        test_dir = Path("test_project_with_error")
+
+        # Ensure the directory doesn't exist from previous runs
+        if test_dir.exists():
+            shutil.rmtree(test_dir)
+
+        # Simulate an error by creating a directory that will cause a failure
+        os.makedirs(test_dir / "src")
+
+        result = self.run_cli("init", str(test_dir))
+        self.assertNotEqual(result.returncode, 0)
+        self.assertFalse(test_dir.exists())  # Directory should be removed on rollback
+
+        # Clean up
+        if test_dir.exists():
+            shutil.rmtree(test_dir)
+
 
 if __name__ == "__main__":
     unittest.main()
