Merge pull request #20 from AgentOps-AI/fix-tools-wizard

Fix tools wizard

Author: bboynton97

agentstack/cli/agentstack_data.py

@@ -63,7 +63,8 @@ def to_json(self):
 
 class FrameworkData:
     def __init__(self,
-                 name: Optional[Literal["crewai"]] = None
+                 # name: Optional[Literal["crewai"]] = None
+                 name: str = None  # TODO: better framework handling, Literal or Enum
                  ):
         self.name = name
 
@@ -80,7 +81,8 @@ class CookiecutterData:
     def __init__(self,
                  project_metadata: ProjectMetadata,
                  structure: ProjectStructure,
-                 framework: Literal["crewai"],
+                 # framework: Literal["crewai"],
+                 framework: str,
                  ):
         self.project_metadata = project_metadata
         self.framework = framework

agentstack/cli/cli.py

@@ -8,12 +8,13 @@
 import inquirer
 import os
 import webbrowser
-import subprocess
 import importlib.resources
 from cookiecutter.main import cookiecutter
 
 from .agentstack_data import FrameworkData, ProjectMetadata, ProjectStructure, CookiecutterData
 from agentstack.logger import log
+from .. import generation
+from ..utils import open_json_file, term_color
 
 
 def init_project_builder(slug_name: Optional[str] = None, skip_wizard: bool = False):
@@ -26,27 +27,29 @@ def init_project_builder(slug_name: Optional[str] = None, skip_wizard: bool = Fa
             "license": "MIT"
         }
 
-        stack = {
-            "framework": "CrewAI"  # TODO: if --no-wizard, require a framework flag
-        }
+        framework = "CrewAI"  # TODO: if --no-wizard, require a framework flag
 
         design = {
             'agents': [],
             'tasks': []
         }
+
+        tools = []
     else:
         welcome_message()
         project_details = ask_project_details(slug_name)
         welcome_message()
-        stack = ask_stack()
+        framework = ask_framework()
         design = ask_design()
+        tools = ask_tools()
 
     log.debug(
         f"project_details: {project_details}"
-        f"stack: {stack}"
+        f"framework: {framework}"
         f"design: {design}"
     )
-    insert_template(project_details, stack, design)
+    insert_template(project_details, framework, design)
+    add_tools(tools, project_details['name'])
 
 
 def welcome_message():
@@ -62,100 +65,37 @@ def welcome_message():
     print(border)
 
 
-def ask_stack():
-    framework = inquirer.prompt(
-        [
-            inquirer.List(
-                "framework",
-                message="What agent framework do you want to use?",
-                choices=["CrewAI", "Autogen", "LiteLLM", "Learn what these are (link)"],
-            )
-        ]
+def ask_framework() -> str:
+    framework = inquirer.list_input(
+        message="What agent framework do you want to use?",
+        choices=["CrewAI", "Autogen", "LiteLLM", "Learn what these are (link)"],
     )
-    if framework["framework"] == "Learn what these are (link)":
+
+    if framework == "Learn what these are (link)":
         webbrowser.open("https://youtu.be/xvFZjo5PgG0")
-        framework = inquirer.prompt(
-            [
-                inquirer.List(
-                    "framework",
-                    message="What agent framework do you want to use?",
-                    choices=["CrewAI", "Autogen", "LiteLLM"],
-                )
-            ]
+        framework = inquirer.list_input(
+            message="What agent framework do you want to use?",
+            choices=["CrewAI", "Autogen", "LiteLLM"],
         )
 
-    while framework["framework"] in ['Autogen', 'LiteLLM']:
-        print(f"{framework['framework']} support coming soon!!")
-        framework = inquirer.prompt(
-            [
-                inquirer.List(
-                    "framework",
-                    message="What agent framework do you want to use?",
-                    choices=["CrewAI", "Autogen", "LiteLLM"],
-                )
-            ]
+    while framework in ['Autogen', 'LiteLLM']:
+        print(f"{framework} support coming soon!!")
+        framework = inquirer.list_input(
+            message="What agent framework do you want to use?",
+            choices=["CrewAI", "Autogen", "LiteLLM"],
         )
 
     print("Congrats! Your project is ready to go! Quickly add features now or skip to do it later.\n\n")
 
-    # TODO: add wizard tool selection back in
-    # use_tools = inquirer.prompt(
-    #     [
-    #         inquirer.Confirm(
-    #             "use_tools",
-    #             message="Do you want to add browsing and RAG tools now? (you can do this later with `agentstack tools add <tool_name>`)",
-    #         )
-    #     ]
-    # )
-
-    use_tools = {'use_tools': False}
-
-    # TODO: dynamically load tools #4
-    browsing_tools = {}
-    rag_tools = {}
-    if use_tools["use_tools"]:
-        browsing_tools = inquirer.prompt(
-            [
-                inquirer.Checkbox(
-                    "browsing_tools",
-                    message="Select browsing tools",
-                    choices=[
-                        "browserbasehq",
-                        "firecrawl",
-                        "MultiOn_AI",
-                        "Crawl4AI",
-                    ],
-                )
-            ]
-        )
-
-        rag_tools = inquirer.prompt(
-            [
-                inquirer.Checkbox(
-                    "rag",
-                    message="RAG/document loading",
-                    choices=[
-                        "Mem0ai",
-                        "llama_index",
-                    ],
-                )
-            ]
-        )
-
-    return {**framework, **browsing_tools, **rag_tools}
+    return framework
 
 
 def ask_design() -> dict:
-    use_wizard = inquirer.prompt(
-        [
-            inquirer.Confirm(
-                "use_wizard",
-                message="Would you like to use the CLI wizard to set up agents and tasks?",
-            )
-        ]
+    use_wizard = inquirer.confirm(
+        message="Would you like to use the CLI wizard to set up agents and tasks?",
     )
 
-    if not use_wizard['use_wizard']:
+    if not use_wizard:
         return {
             'agents': [],
             'tasks': []
@@ -237,6 +177,47 @@ def ask_design() -> dict:
     return {'tasks': tasks, 'agents': agents}
 
 
+def ask_tools() -> list:
+    use_tools = inquirer.confirm(
+        message="Do you want to add agent tools now? (you can do this later with `agentstack tools add <tool_name>`)",
+    )
+
+    if not use_tools:
+        return []
+
+    tools_to_add = []
+
+    adding_tools = True
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    tools_json_path = os.path.join(script_dir, '..', 'tools', 'tools.json')
+
+    # Load the JSON data
+    tools_data = open_json_file(tools_json_path)
+
+    while adding_tools:
+
+        tool_type = inquirer.list_input(
+            message="What category tool do you want to add?",
+            choices=list(tools_data.keys()) + ["~~ Stop adding tools ~~"]
+        )
+
+        tools_in_cat = [f"{t['name']} - {t['url']}" for t in tools_data[tool_type] if t not in tools_to_add]
+        tool_selection = inquirer.list_input(
+            message="Select your tool",
+            choices=tools_in_cat
+        )
+
+        tools_to_add.append(tool_selection.split(' - ')[0])
+
+        print("Adding tools:")
+        for t in tools_to_add:
+            print(f'  - {t}')
+        print('')
+        adding_tools = inquirer.confirm("Add another tool?")
+
+    return tools_to_add
+
+
 def ask_project_details(slug_name: Optional[str] = None) -> dict:
     questions = [
         inquirer.Text("name", message="What's the name of your project (snake_case)", default=slug_name or ''),
@@ -258,8 +239,8 @@ def ask_project_details(slug_name: Optional[str] = None) -> dict:
     return inquirer.prompt(questions)
 
 
-def insert_template(project_details: dict, stack: dict, design: dict):
-    framework = FrameworkData(stack["framework"].lower())
+def insert_template(project_details: dict, framework_name: str, design: dict):
+    framework = FrameworkData(framework_name.lower())
     project_metadata = ProjectMetadata(project_name=project_details["name"],
                                        description=project_details["description"],
                                        author_name=project_details["author"],
@@ -273,7 +254,7 @@ def insert_template(project_details: dict, stack: dict, design: dict):
 
     cookiecutter_data = CookiecutterData(project_metadata=project_metadata,
                                          structure=project_structure,
-                                         framework=stack["framework"].lower())
+                                         framework=framework_name.lower())
 
     with importlib.resources.path(f'agentstack.templates', str(framework.name)) as template_path:
         with open(f"{template_path}/cookiecutter.json", "w") as json_file:
@@ -283,6 +264,11 @@ def insert_template(project_details: dict, stack: dict, design: dict):
         shutil.copy(
             f'{template_path}/{"{{cookiecutter.project_metadata.project_slug}}"}/.env.example',
             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)
 
     # TODO: inits a git repo in the directory the command was run in
@@ -310,15 +296,19 @@ def insert_template(project_details: dict, stack: dict, design: dict):
     )
 
 
+def add_tools(tools: list, project_name: str):
+    for tool in tools:
+        generation.add_tool(tool, project_name)
+
+
 def list_tools():
     try:
         # Determine the path to the tools.json file
         script_dir = os.path.dirname(os.path.abspath(__file__))
         tools_json_path = os.path.join(script_dir, '..', 'tools', 'tools.json')
 
         # Load the JSON data
-        with open(tools_json_path, 'r') as f:
-            tools_data = json.load(f)
+        tools_data = open_json_file(tools_json_path)
 
         # Display the tools
         print("\n\nAvailable AgentStack Tools:")

agentstack/generation/tool_generation.py

@@ -1,43 +1,49 @@
 import sys
+from typing import Optional
+
 from .gen_utils import insert_code_after_tag
 from ..utils import snake_to_camel, open_json_file, get_framework
 import os
 import shutil
 import fileinput
 
 
-def add_tool(tool_name: str):
+def add_tool(tool_name: str, path: Optional[str] = None):
     script_dir = os.path.dirname(os.path.abspath(__file__))
     tools = open_json_file(os.path.join(script_dir, '..', 'tools', 'tools.json'))
-    framework = get_framework()
+    framework = get_framework(path)
     assert_tool_exists(tool_name, tools)
 
     tool_data = open_json_file(os.path.join(script_dir, '..', 'tools', f'{tool_name}.json'))
     tool_file_route = os.path.join(script_dir, '..', 'templates', framework, 'tools', f'{tool_name}.py')
 
     os.system(tool_data['package'])  # Install package
-    shutil.copy(tool_file_route, f'src/tools/{tool_name}.py')  # Move tool from package to project
-    add_tool_to_tools_init(tool_data)  # Export tool from tools dir
-    add_tool_to_agent_definition(framework, tool_data)
-    insert_code_after_tag('.env', '# Tools', [tool_data['env']], next_line=True)  # Add env var
-    insert_code_after_tag('.env.example', '# Tools', [tool_data['env']], next_line=True)  # Add env var
+    shutil.copy(tool_file_route, f'{path or ""}/src/tools/{tool_name}.py')  # Move tool from package to project
+    add_tool_to_tools_init(tool_data, path)  # Export tool from tools dir
+    add_tool_to_agent_definition(framework, tool_data, path)
+    insert_code_after_tag(f'{path}/.env', '# Tools', [tool_data['env']], next_line=True)  # Add env var
+    insert_code_after_tag(f'{path}/.env.example', '# Tools', [tool_data['env']], next_line=True)  # Add env var
 
     print(f'\033[92m🔨 Tool {tool_name} added to agentstack project successfully\033[0m')
 
-def add_tool_to_tools_init(tool_data: dict):
-    file_path = 'src/tools/__init__.py'
+
+def add_tool_to_tools_init(tool_data: dict, path: Optional[str] = None):
+    file_path = f'{path or ""}/src/tools/__init__.py'
     tag = '# tool import'
     code_to_insert = [
         f"from {tool_data['name']} import {', '.join([tool_name for tool_name in tool_data['tools']])}"
     ]
     insert_code_after_tag(file_path, tag, code_to_insert, next_line=True)
 
 
-def add_tool_to_agent_definition(framework: str, tool_data: dict):
+def add_tool_to_agent_definition(framework: str, tool_data: dict, path: Optional[str] = None):
     filename = ''
     if framework == 'crewai':
         filename = 'src/crew.py'
 
+    if path:
+        filename = f'{path}/{filename}'
+
     with fileinput.input(files=filename, inplace=True) as f:
         for line in f:
             print(line.replace('tools=[', f'tools=[tools.{", tools.".join([tool_name for tool_name in tool_data["tools"]])}, '), end='')

agentstack/utils.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import toml
 import os
 import sys
@@ -23,9 +25,12 @@ def verify_agentstack_project():
         sys.exit(1)
 
 
-def get_framework() -> str:
+def get_framework(path: Optional[str] = None) -> str:
     try:
-        with open('agentstack.json', 'r') as f:
+        file_path = 'agentstack.json'
+        if path is not None:
+            file_path = path + '/' + file_path
+        with open(file_path, 'r') as f:
             data = json.load(f)
             framework = data.get('framework')
 
@@ -47,7 +52,7 @@ def snake_to_camel(s):
     return ''.join(word.title() for word in s.split('_'))
 
 
-def open_json_file(path):
+def open_json_file(path) -> dict:
     with open(path, 'r') as f:
         data = json.load(f)
     return data
@@ -56,3 +61,12 @@ def open_json_file(path):
 def clean_input(input_string):
     special_char_pattern = re.compile(r'[^a-zA-Z0-9\s_]')
     return re.sub(special_char_pattern, '', input_string).lower().replace(' ', '_').replace('-', '_')
+
+
+def term_color(text: str, color: str) -> str:
+    if color is 'red':
+        return "\033[91m{}\033[00m".format(text)
+    if color is 'green':
+        return "\033[92m{}\033[00m".format(text)
+    else:
+        return text