Skip to content

Commit

Permalink
Merge pull request #64 from tcdent/issue-58
Browse files Browse the repository at this point in the history 
Validate tool config json file on load
  • Loading branch information
@bboynton97
bboynton97 authored Nov 28, 2024
2 parents e55de28 + 3f25983 commit 947865d
Show file tree
Hide file tree
Showing 2 changed files with 125 additions and 105 deletions.
227 changes: 123 additions & 104 deletions agentstack/generation/tool_generation.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import importlib.resources
from pathlib import Path
import json
import sys
from typing import Optional
from typing import Optional, List
from pydantic import BaseModel, ValidationError

from .gen_utils import insert_code_after_tag, string_in_file
from ..utils import open_json_file, get_framework, term_color
Expand All @@ -10,135 +12,152 @@
import fileinput

TOOL_INIT_FILENAME = "src/tools/__init__.py"
TOOLS_DATA_PATH: Path = importlib.resources.files('agentstack.tools') / 'tools.json'
AGENTSTACK_JSON_FILENAME = "agentstack.json"
FRAMEWORK_FILENAMES: dict[str, str] = {
'crewai': 'src/crew.py',
}

def get_framework_filename(framework: str, path: str = ''):
try:
return FRAMEWORK_FILENAMES[framework]
except KeyError:
print(term_color(f'Unknown framework: {framework}', 'red'))
sys.exit(1)

def assert_tool_exists(name: str):
tools_data = open_json_file(TOOLS_DATA_PATH)
for category, tools in tools_data.items():
for tool_dict in tools:
if tool_dict['name'] == name:
return
print(term_color(f'No known agentstack tool: {name}', 'red'))
sys.exit(1)

class ToolConfig(BaseModel):
name: str
tools: list[str]
tools_bundled: bool = False
cta: Optional[str] = None
env: Optional[str] = None
packages: Optional[List[str]] = None
post_install: Optional[str] = None
post_remove: Optional[str] = None

@classmethod
def from_tool_name(cls, name: str) -> 'ToolConfig':
assert_tool_exists(name)
return cls.from_json(importlib.resources.files('agentstack.tools') / f'{name}.json')

@classmethod
def from_json(cls, path: Path) -> 'ToolConfig':
data = open_json_file(path)
try:
return cls(**data)
except ValidationError as e:
print(term_color(f"Error validating tool config JSON: \n{path}", 'red'))
for error in e.errors():
print(f"{' '.join(error['loc'])}: {error['msg']}")
sys.exit(1)

def get_import_statement(self) -> str:
return f"from .{self.name}_tool import {', '.join(self.tools)}"

def add_tool(tool_name: str, path: Optional[str] = None):
if path:
path = path.endswith('/') and path or path + '/'
else:
path = './'
with importlib.resources.path(f'agentstack.tools', 'tools.json') as tools_data_path:
tools = open_json_file(tools_data_path)
framework = get_framework(path)
assert_tool_exists(tool_name, tools)
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')

if tool_name in agentstack_json.get('tools', []):
print(term_color(f'Tool {tool_name} is already installed', 'red'))
sys.exit(1)

with importlib.resources.path(f'agentstack.tools', f"{tool_name}.json") as tool_data_path:
tool_data = open_json_file(tool_data_path)

with importlib.resources.path(f'agentstack.templates.{framework}.tools', f"{tool_name}_tool.py") as tool_file_path:
if tool_data.get('packages'):
os.system(f"poetry add {' '.join(tool_data['packages'])}") # Install packages
shutil.copy(tool_file_path, f'{path}src/tools/{tool_name}_tool.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) # Add tool to agent definition
if tool_data.get('env'): # if the env vars aren't in the .env files, add them
first_var_name = tool_data['env'].split('=')[0]
if not string_in_file(f'{path}.env', first_var_name):
insert_code_after_tag(f'{path}.env', '# Tools', [tool_data['env']], next_line=True) # Add env var
if not string_in_file(f'{path}.env.example', first_var_name):
insert_code_after_tag(f'{path}.env.example', '# Tools', [tool_data['env']], next_line=True) # Add env var
if tool_data.get('post_install'):
os.system(tool_data['post_install'])
if not agentstack_json.get('tools'):
agentstack_json['tools'] = []
agentstack_json['tools'].append(tool_name)

with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
json.dump(agentstack_json, f, indent=4)

print(term_color(f'🔨 Tool {tool_name} added to agentstack project successfully', 'green'))
if tool_data.get('cta'):
print(term_color(f'🪩 {tool_data["cta"]}', 'blue'))


framework = get_framework(path)
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')

if tool_name in agentstack_json.get('tools', []):
print(term_color(f'Tool {tool_name} is already installed', 'red'))
sys.exit(1)

tool_data = ToolConfig.from_tool_name(tool_name)
tool_file_path = importlib.resources.files(f'agentstack.templates.{framework}.tools') / f'{tool_name}_tool.py'
if tool_data.packages:
os.system(f"poetry add {' '.join(tool_data.packages)}") # Install packages
shutil.copy(tool_file_path, f'{path}src/tools/{tool_name}_tool.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) # Add tool to agent definition
if tool_data.env: # if the env vars aren't in the .env files, add them
first_var_name = tool_data.env.split('=')[0]
if not string_in_file(f'{path}.env', first_var_name):
insert_code_after_tag(f'{path}.env', '# Tools', [tool_data.env], next_line=True) # Add env var
if not string_in_file(f'{path}.env.example', first_var_name):
insert_code_after_tag(f'{path}.env.example', '# Tools', [tool_data.env], next_line=True) # Add env var

if tool_data.post_install:
os.system(tool_data.post_install)

if not agentstack_json.get('tools'):
agentstack_json['tools'] = []
agentstack_json['tools'].append(tool_name)

with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
json.dump(agentstack_json, f, indent=4)

print(term_color(f'🔨 Tool {tool_name} added to agentstack project successfully', 'green'))
if tool_data.cta:
print(term_color(f'🪩 {tool_data.cta}', 'blue'))

def remove_tool(tool_name: str, path: Optional[str] = None):
if path:
path = path.endswith('/') and path or path + '/'
else:
path = './'
with importlib.resources.path(f'agentstack.tools', 'tools.json') as tools_data_path:
tools = open_json_file(tools_data_path)
framework = get_framework()
assert_tool_exists(tool_name, tools)
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')

if not tool_name in agentstack_json.get('tools', []):
print(term_color(f'Tool {tool_name} is not installed', 'red'))
sys.exit(1)

with importlib.resources.path(f'agentstack.tools', f"{tool_name}.json") as tool_data_path:
tool_data = open_json_file(tool_data_path)
if tool_data.get('packages'):
os.system(f"poetry remove {' '.join(tool_data['packages'])}") # Uninstall packages
os.remove(f'{path}src/tools/{tool_name}_tool.py')
remove_tool_from_tools_init(tool_data, path)
remove_tool_from_agent_definition(framework, tool_data, path)
if tool_data.get('post_remove'):
os.system(tool_data['post_remove'])
# We don't remove the .env variables to preserve user data.

agentstack_json['tools'].remove(tool_name)
with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
json.dump(agentstack_json, f, indent=4)

print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), term_color('from agentstack project successfully', 'green'))


def _format_tool_import_statement(tool_data: dict):
return f"from .{tool_data['name']}_tool import {', '.join([tool_name for tool_name in tool_data['tools']])}"


def add_tool_to_tools_init(tool_data: dict, path: str = ''):

framework = get_framework()
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')

if not tool_name in agentstack_json.get('tools', []):
print(term_color(f'Tool {tool_name} is not installed', 'red'))
sys.exit(1)

tool_data = ToolConfig.from_tool_name(tool_name)
if tool_data.packages:
os.system(f"poetry remove {' '.join(tool_data.packages)}") # Uninstall packages
try:
os.remove(f'{path}src/tools/{tool_name}_tool.py')
except FileNotFoundError:
print(f'"src/tools/{tool_name}_tool.py" not found')
remove_tool_from_tools_init(tool_data, path)
remove_tool_from_agent_definition(framework, tool_data, path)
if tool_data.post_remove:
os.system(tool_data.post_remove)
# We don't remove the .env variables to preserve user data.

agentstack_json['tools'].remove(tool_name)
with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
json.dump(agentstack_json, f, indent=4)

print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), term_color('from agentstack project successfully', 'green'))

def add_tool_to_tools_init(tool_data: ToolConfig, path: str = ''):
file_path = f'{path}{TOOL_INIT_FILENAME}'
tag = '# tool import'
code_to_insert = [_format_tool_import_statement(tool_data), ]
code_to_insert = [tool_data.get_import_statement(), ]
insert_code_after_tag(file_path, tag, code_to_insert, next_line=True)


def remove_tool_from_tools_init(tool_data: dict, path: str = ''):
def remove_tool_from_tools_init(tool_data: ToolConfig, path: str = ''):
"""Search for the import statement in the init and remove it."""
file_path = f'{path}{TOOL_INIT_FILENAME}'
import_statement = _format_tool_import_statement(tool_data)
import_statement = tool_data.get_import_statement()
with fileinput.input(files=file_path, inplace=True) as f:
for line in f:
if line.strip() != import_statement:
print(line, end='')


def _framework_filename(framework: str, path: str = ''):
if framework == 'crewai':
return f'{path}src/crew.py'

print(term_color(f'Unknown framework: {framework}', 'red'))
sys.exit(1)


def add_tool_to_agent_definition(framework: str, tool_data: dict, path: str = ''):
filename = _framework_filename(framework, path)
with fileinput.input(files=filename, inplace=True) as f:
def add_tool_to_agent_definition(framework: str, tool_data: ToolConfig, path: str = ''):
with fileinput.input(files=get_framework_filename(framework, path), inplace=True) as f:
for line in f:
print(line.replace('tools=[', f'tools=[{"*" if tool_data.get("tools_bundled") else ""}tools.{", tools.".join([tool_name for tool_name in tool_data["tools"]])}, '), end='')

print(line.replace('tools=[', f'tools=[{"*" if tool_data.tools_bundled else ""}tools.{", tools.".join(tool_data.tools)}, '), end='')

def remove_tool_from_agent_definition(framework: str, tool_data: dict, path: str = ''):
filename = _framework_filename(framework, path)
with fileinput.input(files=filename, inplace=True) as f:
def remove_tool_from_agent_definition(framework: str, tool_data: ToolConfig, path: str = ''):
with fileinput.input(files=get_framework_filename(framework, path), inplace=True) as f:
for line in f:
print(line.replace(f'{", ".join([f"tools.{tool_name}" for tool_name in tool_data["tools"]])}, ', ''), end='')


def assert_tool_exists(tool_name: str, tools: dict):
for cat in tools.keys():
for tool_dict in tools[cat]:
if tool_dict['name'] == tool_name:
return

print(term_color(f'No known agentstack tool: {tool_name}', 'red'))
sys.exit(1)
print(line.replace(f'{", ".join([f"tools.{name}" for name in tool_data.tools])}, ', ''), end='')

3 changes: 2 additions & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ dependencies = [
"toml>=0.10.2",
"ruamel.yaml.base>=0.3.2",
"cookiecutter==2.6.0",
"psutil==5.9.0"
"psutil==5.9.0",
"pydantic>=2.10",
]

[tool.setuptools.package-data]
Expand Down

0 comments on commit 947865d

Please sign in to comment.