tool gen from wizard

Author: bboynton97

agentstack/cli/cli.py

@@ -13,7 +13,8 @@
 
 from .agentstack_data import FrameworkData, ProjectMetadata, ProjectStructure, CookiecutterData
 from agentstack.logger import log
-from ..utils import open_json_file
+from .. import generation
+from ..utils import open_json_file, term_color
 
 
 def init_project_builder(slug_name: Optional[str] = None, skip_wizard: bool = False):
@@ -32,6 +33,8 @@ def init_project_builder(slug_name: Optional[str] = None, skip_wizard: bool = Fa
             'agents': [],
             'tasks': []
         }
+
+        tools = []
     else:
         welcome_message()
         project_details = ask_project_details(slug_name)
@@ -46,6 +49,7 @@ def init_project_builder(slug_name: Optional[str] = None, skip_wizard: bool = Fa
         f"design: {design}"
     )
     insert_template(project_details, framework, design)
+    add_tools(tools, project_details['name'])
 
 
 def welcome_message():
@@ -260,6 +264,11 @@ def insert_template(project_details: dict, framework_name: str, 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
@@ -287,6 +296,11 @@ def insert_template(project_details: dict, framework_name: str, 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

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')
 
@@ -56,3 +61,12 @@ def open_json_file(path) -> dict:
 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