Skip to content

Commit 947865d

Browse files
bboynton97bboynton97
authored
Merge pull request #64 from tcdent/issue-58
Validate tool config json file on load
2 parents e55de28 + 3f25983 commit 947865d

2 files changed

Lines changed: 125 additions & 105 deletions

File tree

‎agentstack/generation/tool_generation.py‎

Lines changed: 123 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
import importlib.resources
2+
from pathlib import Path
23
import json
34
import sys
4-
from typing import Optional
5+
from typing import Optional, List
6+
from pydantic import BaseModel, ValidationError
57

68
from .gen_utils import insert_code_after_tag, string_in_file
79
from ..utils import open_json_file, get_framework, term_color
@@ -10,135 +12,152 @@
1012
import fileinput
1113

1214
TOOL_INIT_FILENAME = "src/tools/__init__.py"
15+
TOOLS_DATA_PATH: Path = importlib.resources.files('agentstack.tools') / 'tools.json'
1316
AGENTSTACK_JSON_FILENAME = "agentstack.json"
17+
FRAMEWORK_FILENAMES: dict[str, str] = {
18+
'crewai': 'src/crew.py',
19+
}
20+
21+
def get_framework_filename(framework: str, path: str = ''):
22+
try:
23+
return FRAMEWORK_FILENAMES[framework]
24+
except KeyError:
25+
print(term_color(f'Unknown framework: {framework}', 'red'))
26+
sys.exit(1)
27+
28+
def assert_tool_exists(name: str):
29+
tools_data = open_json_file(TOOLS_DATA_PATH)
30+
for category, tools in tools_data.items():
31+
for tool_dict in tools:
32+
if tool_dict['name'] == name:
33+
return
34+
print(term_color(f'No known agentstack tool: {name}', 'red'))
35+
sys.exit(1)
1436

37+
class ToolConfig(BaseModel):
38+
name: str
39+
tools: list[str]
40+
tools_bundled: bool = False
41+
cta: Optional[str] = None
42+
env: Optional[str] = None
43+
packages: Optional[List[str]] = None
44+
post_install: Optional[str] = None
45+
post_remove: Optional[str] = None
46+
47+
@classmethod
48+
def from_tool_name(cls, name: str) -> 'ToolConfig':
49+
assert_tool_exists(name)
50+
return cls.from_json(importlib.resources.files('agentstack.tools') / f'{name}.json')
51+
52+
@classmethod
53+
def from_json(cls, path: Path) -> 'ToolConfig':
54+
data = open_json_file(path)
55+
try:
56+
return cls(**data)
57+
except ValidationError as e:
58+
print(term_color(f"Error validating tool config JSON: \n{path}", 'red'))
59+
for error in e.errors():
60+
print(f"{' '.join(error['loc'])}: {error['msg']}")
61+
sys.exit(1)
62+
63+
def get_import_statement(self) -> str:
64+
return f"from .{self.name}_tool import {', '.join(self.tools)}"
1565

1666
def add_tool(tool_name: str, path: Optional[str] = None):
1767
if path:
1868
path = path.endswith('/') and path or path + '/'
1969
else:
2070
path = './'
21-
with importlib.resources.path(f'agentstack.tools', 'tools.json') as tools_data_path:
22-
tools = open_json_file(tools_data_path)
23-
framework = get_framework(path)
24-
assert_tool_exists(tool_name, tools)
25-
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')
26-
27-
if tool_name in agentstack_json.get('tools', []):
28-
print(term_color(f'Tool {tool_name} is already installed', 'red'))
29-
sys.exit(1)
30-
31-
with importlib.resources.path(f'agentstack.tools', f"{tool_name}.json") as tool_data_path:
32-
tool_data = open_json_file(tool_data_path)
33-
34-
with importlib.resources.path(f'agentstack.templates.{framework}.tools', f"{tool_name}_tool.py") as tool_file_path:
35-
if tool_data.get('packages'):
36-
os.system(f"poetry add {' '.join(tool_data['packages'])}") # Install packages
37-
shutil.copy(tool_file_path, f'{path}src/tools/{tool_name}_tool.py') # Move tool from package to project
38-
add_tool_to_tools_init(tool_data, path) # Export tool from tools dir
39-
add_tool_to_agent_definition(framework, tool_data, path) # Add tool to agent definition
40-
if tool_data.get('env'): # if the env vars aren't in the .env files, add them
41-
first_var_name = tool_data['env'].split('=')[0]
42-
if not string_in_file(f'{path}.env', first_var_name):
43-
insert_code_after_tag(f'{path}.env', '# Tools', [tool_data['env']], next_line=True) # Add env var
44-
if not string_in_file(f'{path}.env.example', first_var_name):
45-
insert_code_after_tag(f'{path}.env.example', '# Tools', [tool_data['env']], next_line=True) # Add env var
46-
if tool_data.get('post_install'):
47-
os.system(tool_data['post_install'])
48-
if not agentstack_json.get('tools'):
49-
agentstack_json['tools'] = []
50-
agentstack_json['tools'].append(tool_name)
51-
52-
with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
53-
json.dump(agentstack_json, f, indent=4)
54-
55-
print(term_color(f'🔨 Tool {tool_name} added to agentstack project successfully', 'green'))
56-
if tool_data.get('cta'):
57-
print(term_color(f'🪩 {tool_data["cta"]}', 'blue'))
58-
71+
72+
framework = get_framework(path)
73+
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')
74+
75+
if tool_name in agentstack_json.get('tools', []):
76+
print(term_color(f'Tool {tool_name} is already installed', 'red'))
77+
sys.exit(1)
78+
79+
tool_data = ToolConfig.from_tool_name(tool_name)
80+
tool_file_path = importlib.resources.files(f'agentstack.templates.{framework}.tools') / f'{tool_name}_tool.py'
81+
if tool_data.packages:
82+
os.system(f"poetry add {' '.join(tool_data.packages)}") # Install packages
83+
shutil.copy(tool_file_path, f'{path}src/tools/{tool_name}_tool.py') # Move tool from package to project
84+
add_tool_to_tools_init(tool_data, path) # Export tool from tools dir
85+
add_tool_to_agent_definition(framework, tool_data, path) # Add tool to agent definition
86+
if tool_data.env: # if the env vars aren't in the .env files, add them
87+
first_var_name = tool_data.env.split('=')[0]
88+
if not string_in_file(f'{path}.env', first_var_name):
89+
insert_code_after_tag(f'{path}.env', '# Tools', [tool_data.env], next_line=True) # Add env var
90+
if not string_in_file(f'{path}.env.example', first_var_name):
91+
insert_code_after_tag(f'{path}.env.example', '# Tools', [tool_data.env], next_line=True) # Add env var
92+
93+
if tool_data.post_install:
94+
os.system(tool_data.post_install)
95+
96+
if not agentstack_json.get('tools'):
97+
agentstack_json['tools'] = []
98+
agentstack_json['tools'].append(tool_name)
99+
100+
with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
101+
json.dump(agentstack_json, f, indent=4)
102+
103+
print(term_color(f'🔨 Tool {tool_name} added to agentstack project successfully', 'green'))
104+
if tool_data.cta:
105+
print(term_color(f'🪩 {tool_data.cta}', 'blue'))
59106

60107
def remove_tool(tool_name: str, path: Optional[str] = None):
61108
if path:
62109
path = path.endswith('/') and path or path + '/'
63110
else:
64111
path = './'
65-
with importlib.resources.path(f'agentstack.tools', 'tools.json') as tools_data_path:
66-
tools = open_json_file(tools_data_path)
67-
framework = get_framework()
68-
assert_tool_exists(tool_name, tools)
69-
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')
70-
71-
if not tool_name in agentstack_json.get('tools', []):
72-
print(term_color(f'Tool {tool_name} is not installed', 'red'))
73-
sys.exit(1)
74-
75-
with importlib.resources.path(f'agentstack.tools', f"{tool_name}.json") as tool_data_path:
76-
tool_data = open_json_file(tool_data_path)
77-
if tool_data.get('packages'):
78-
os.system(f"poetry remove {' '.join(tool_data['packages'])}") # Uninstall packages
79-
os.remove(f'{path}src/tools/{tool_name}_tool.py')
80-
remove_tool_from_tools_init(tool_data, path)
81-
remove_tool_from_agent_definition(framework, tool_data, path)
82-
if tool_data.get('post_remove'):
83-
os.system(tool_data['post_remove'])
84-
# We don't remove the .env variables to preserve user data.
85-
86-
agentstack_json['tools'].remove(tool_name)
87-
with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
88-
json.dump(agentstack_json, f, indent=4)
89-
90-
print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), term_color('from agentstack project successfully', 'green'))
91-
92-
93-
def _format_tool_import_statement(tool_data: dict):
94-
return f"from .{tool_data['name']}_tool import {', '.join([tool_name for tool_name in tool_data['tools']])}"
95-
96-
97-
def add_tool_to_tools_init(tool_data: dict, path: str = ''):
112+
113+
framework = get_framework()
114+
agentstack_json = open_json_file(f'{path}{AGENTSTACK_JSON_FILENAME}')
115+
116+
if not tool_name in agentstack_json.get('tools', []):
117+
print(term_color(f'Tool {tool_name} is not installed', 'red'))
118+
sys.exit(1)
119+
120+
tool_data = ToolConfig.from_tool_name(tool_name)
121+
if tool_data.packages:
122+
os.system(f"poetry remove {' '.join(tool_data.packages)}") # Uninstall packages
123+
try:
124+
os.remove(f'{path}src/tools/{tool_name}_tool.py')
125+
except FileNotFoundError:
126+
print(f'"src/tools/{tool_name}_tool.py" not found')
127+
remove_tool_from_tools_init(tool_data, path)
128+
remove_tool_from_agent_definition(framework, tool_data, path)
129+
if tool_data.post_remove:
130+
os.system(tool_data.post_remove)
131+
# We don't remove the .env variables to preserve user data.
132+
133+
agentstack_json['tools'].remove(tool_name)
134+
with open(f'{path}{AGENTSTACK_JSON_FILENAME}', 'w') as f:
135+
json.dump(agentstack_json, f, indent=4)
136+
137+
print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), term_color('from agentstack project successfully', 'green'))
138+
139+
def add_tool_to_tools_init(tool_data: ToolConfig, path: str = ''):
98140
file_path = f'{path}{TOOL_INIT_FILENAME}'
99141
tag = '# tool import'
100-
code_to_insert = [_format_tool_import_statement(tool_data), ]
142+
code_to_insert = [tool_data.get_import_statement(), ]
101143
insert_code_after_tag(file_path, tag, code_to_insert, next_line=True)
102144

103-
104-
def remove_tool_from_tools_init(tool_data: dict, path: str = ''):
145+
def remove_tool_from_tools_init(tool_data: ToolConfig, path: str = ''):
105146
"""Search for the import statement in the init and remove it."""
106147
file_path = f'{path}{TOOL_INIT_FILENAME}'
107-
import_statement = _format_tool_import_statement(tool_data)
148+
import_statement = tool_data.get_import_statement()
108149
with fileinput.input(files=file_path, inplace=True) as f:
109150
for line in f:
110151
if line.strip() != import_statement:
111152
print(line, end='')
112153

113-
114-
def _framework_filename(framework: str, path: str = ''):
115-
if framework == 'crewai':
116-
return f'{path}src/crew.py'
117-
118-
print(term_color(f'Unknown framework: {framework}', 'red'))
119-
sys.exit(1)
120-
121-
122-
def add_tool_to_agent_definition(framework: str, tool_data: dict, path: str = ''):
123-
filename = _framework_filename(framework, path)
124-
with fileinput.input(files=filename, inplace=True) as f:
154+
def add_tool_to_agent_definition(framework: str, tool_data: ToolConfig, path: str = ''):
155+
with fileinput.input(files=get_framework_filename(framework, path), inplace=True) as f:
125156
for line in f:
126-
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='')
127-
157+
print(line.replace('tools=[', f'tools=[{"*" if tool_data.tools_bundled else ""}tools.{", tools.".join(tool_data.tools)}, '), end='')
128158

129-
def remove_tool_from_agent_definition(framework: str, tool_data: dict, path: str = ''):
130-
filename = _framework_filename(framework, path)
131-
with fileinput.input(files=filename, inplace=True) as f:
159+
def remove_tool_from_agent_definition(framework: str, tool_data: ToolConfig, path: str = ''):
160+
with fileinput.input(files=get_framework_filename(framework, path), inplace=True) as f:
132161
for line in f:
133-
print(line.replace(f'{", ".join([f"tools.{tool_name}" for tool_name in tool_data["tools"]])}, ', ''), end='')
134-
135-
136-
def assert_tool_exists(tool_name: str, tools: dict):
137-
for cat in tools.keys():
138-
for tool_dict in tools[cat]:
139-
if tool_dict['name'] == tool_name:
140-
return
141-
142-
print(term_color(f'No known agentstack tool: {tool_name}', 'red'))
143-
sys.exit(1)
162+
print(line.replace(f'{", ".join([f"tools.{name}" for name in tool_data.tools])}, ', ''), end='')
144163

‎pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ dependencies = [
2020
"toml>=0.10.2",
2121
"ruamel.yaml.base>=0.3.2",
2222
"cookiecutter==2.6.0",
23-
"psutil==5.9.0"
23+
"psutil==5.9.0",
24+
"pydantic>=2.10",
2425
]
2526

2627
[tool.setuptools.package-data]

0 commit comments

Comments
 (0)