Skip to content

Commit

Permalink
tool gen from wizard
Browse files Browse the repository at this point in the history
  • Loading branch information
bboynton97 committed Oct 17, 2024
1 parent 0cad2cd commit 9a95d0d
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 13 deletions.
16 changes: 15 additions & 1 deletion agentstack/cli/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Expand All @@ -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():
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
26 changes: 16 additions & 10 deletions agentstack/generation/tool_generation.py
Original file line number Diff line number Diff line change
@@ -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='')
Expand Down
18 changes: 16 additions & 2 deletions agentstack/utils.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Optional

import toml
import os
import sys
Expand All @@ -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')

Expand Down Expand Up @@ -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

0 comments on commit 9a95d0d

Please sign in to comment.