diff --git a/agentstack/generation/__init__.py b/agentstack/generation/__init__.py index 41bbc9c..b9aa094 100644 --- a/agentstack/generation/__init__.py +++ b/agentstack/generation/__init__.py @@ -1,3 +1,3 @@ from .agent_generation import generate_agent from .task_generation import generate_task -from .tool_generation import add_tool +from .tool_generation import add_tool, remove_tool diff --git a/agentstack/generation/gen_utils.py b/agentstack/generation/gen_utils.py index 3faf78c..d0aabfa 100644 --- a/agentstack/generation/gen_utils.py +++ b/agentstack/generation/gen_utils.py @@ -59,3 +59,9 @@ def insert_after_tasks(file_path, code_to_insert): return True return False + +def string_in_file(file_path: str, str_to_match: str) -> bool: + with open(file_path, 'r') as file: + file_content = file.read() + return str_to_match in file_content + diff --git a/agentstack/generation/tool_generation.py b/agentstack/generation/tool_generation.py index c50c48b..f60b3f1 100644 --- a/agentstack/generation/tool_generation.py +++ b/agentstack/generation/tool_generation.py @@ -3,36 +3,52 @@ import sys from typing import Optional -from .gen_utils import insert_code_after_tag +from .gen_utils import insert_code_after_tag, string_in_file from ..utils import open_json_file, get_framework, term_color import os import shutil import fileinput +TOOL_INIT_FILENAME = "src/tools/__init__.py" +AGENTSTACK_JSON_FILENAME = "agentstack.json" + 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: - os.system(tool_data['package']) # Install package - shutil.copy(tool_file_path, f'{path + "/" if path else ""}src/tools/{tool_name}_tool.py') # Move tool from package to project + 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) - insert_code_after_tag(f'{path + "/" if path else ""}.env', '# Tools', [tool_data['env']], next_line=True) # Add env var - insert_code_after_tag(f'{path + "/" if path else ""}.env.example', '# Tools', [tool_data['env']], next_line=True) # Add env var - - agentstack_json = open_json_file(f'{path + "/" if path else ""}agentstack.json') + 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 not agentstack_json.get('tools'): agentstack_json['tools'] = [] agentstack_json['tools'].append(tool_name) - with open(f'{path + "/" if path else ""}agentstack.json', 'w') as f: + 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')) @@ -40,34 +56,86 @@ def add_tool(tool_name: str, path: Optional[str] = None): print(term_color(f'🪩 {tool_data["cta"]}', 'blue')) -def add_tool_to_tools_init(tool_data: dict, path: Optional[str] = None): - file_path = f'{path + "/" if path else ""}src/tools/__init__.py' +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) + # 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 = ''): + file_path = f'{path}{TOOL_INIT_FILENAME}' tag = '# tool import' - code_to_insert = [ - f"from .{tool_data['name']}_tool import {', '.join([tool_name for tool_name in tool_data['tools']])}" - ] + code_to_insert = [_format_tool_import_statement(tool_data), ] insert_code_after_tag(file_path, tag, code_to_insert, next_line=True) -def add_tool_to_agent_definition(framework: str, tool_data: dict, path: Optional[str] = None): - filename = '' +def remove_tool_from_tools_init(tool_data: dict, 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) + 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': - filename = 'src/crew.py' + return f'{path}src/crew.py' - if path: - filename = f'{path}/{filename}' + 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: 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='') +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: + 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(f"\033[31mNo known AgentStack tool: '{tool_name}'\033[0m") + print(term_color(f'No known agentstack tool: {tool_name}', 'red')) sys.exit(1) diff --git a/agentstack/main.py b/agentstack/main.py index 09c7d41..fdc6180 100644 --- a/agentstack/main.py +++ b/agentstack/main.py @@ -62,6 +62,10 @@ def main(): tools_add_parser = tools_subparsers.add_parser('add', aliases=['a'], help='Add a new tool') tools_add_parser.add_argument('name', help='Name of the tool to add') + # 'remove' command under 'tools' + tools_remove_parser = tools_subparsers.add_parser('remove', aliases=['r'], help='Remove a tool') + tools_remove_parser.add_argument('name', help='Name of the tool to remove') + # Parse arguments args = parser.parse_args() @@ -89,6 +93,8 @@ def main(): list_tools() elif args.tools_command in ['add', 'a']: generation.add_tool(args.name) + elif args.tools_command in ['remove', 'r']: + generation.remove_tool(args.name) else: tools_parser.print_help() else: diff --git a/agentstack/tools/agent_connect.json b/agentstack/tools/agent_connect.json index 43c87b4..34922eb 100644 --- a/agentstack/tools/agent_connect.json +++ b/agentstack/tools/agent_connect.json @@ -1,6 +1,6 @@ { "name": "agent-connect", - "package": "poetry add agent-connect", + "packages": ["agent-connect"], "env": "HOST_DOMAIN=...\nHOST_PORT=\"80\"\nHOST_WS_PATH=\"/ws\"\nDID_DOCUMENT_PATH=...\nSSL_CERT_PATH=...\nSSL_KEY_PATH=...", "tools": ["send_message", "receive_message"] } diff --git a/agentstack/tools/browserbase.json b/agentstack/tools/browserbase.json index 4ffb1a0..9e75455 100644 --- a/agentstack/tools/browserbase.json +++ b/agentstack/tools/browserbase.json @@ -1,6 +1,6 @@ { "name": "browserbase", - "package": "poetry add browserbase playwright", + "packages": ["browserbase", "playwright"], "env": "BROWSERBASE_API_KEY=...\nBROWSERBASE_PROJECT_ID=...", "tools": ["browserbase"], "cta": "Create an API key at https://www.browserbase.com/" diff --git a/agentstack/tools/code_interpreter.json b/agentstack/tools/code_interpreter.json index 31271db..a746ff1 100644 --- a/agentstack/tools/code_interpreter.json +++ b/agentstack/tools/code_interpreter.json @@ -1,6 +1,6 @@ { "name": "code_interpreter", - "package": "poetry add crewai-tools", + "packages": [], "env": "", "tools": ["code_interpreter"] } \ No newline at end of file diff --git a/agentstack/tools/composio.json b/agentstack/tools/composio.json index 2a80bc6..c89d699 100644 --- a/agentstack/tools/composio.json +++ b/agentstack/tools/composio.json @@ -1,6 +1,6 @@ { "name": "composio", - "package": "poetry add composio-crewai", + "packages": ["composio-crewai"], "env": "COMPOSIO_API_KEY=...", "tools": ["composio_tools"], "tools_bundled": true, diff --git a/agentstack/tools/directory_search.json b/agentstack/tools/directory_search.json index 5ac6648..59fb371 100644 --- a/agentstack/tools/directory_search.json +++ b/agentstack/tools/directory_search.json @@ -1,6 +1,6 @@ { "name": "dir_search_tool", - "package": "poetry add crewai-tools", + "packages": [], "env": "", "tools": ["dir_search_tool"] } \ No newline at end of file diff --git a/agentstack/tools/exa.json b/agentstack/tools/exa.json index 45690ac..fbcec41 100644 --- a/agentstack/tools/exa.json +++ b/agentstack/tools/exa.json @@ -1,6 +1,6 @@ { "name": "exa", - "package": "poetry add exa_py", + "packages": ["exa_py"], "env": "EXA_API_KEY=...", "tools": ["search_and_contents"], "cta": "Get your Exa API key at https://dashboard.exa.ai/api-keys" diff --git a/agentstack/tools/file_read.json b/agentstack/tools/file_read.json index 04be812..98620c3 100644 --- a/agentstack/tools/file_read.json +++ b/agentstack/tools/file_read.json @@ -1,6 +1,6 @@ { "name": "file_read_tool", - "package": "poetry add crewai-tools", + "packages": [], "env": "", "tools": ["file_read_tool"] } \ No newline at end of file diff --git a/agentstack/tools/firecrawl.json b/agentstack/tools/firecrawl.json index e16d742..2fcdaae 100644 --- a/agentstack/tools/firecrawl.json +++ b/agentstack/tools/firecrawl.json @@ -1,6 +1,6 @@ { "name": "firecrawl", - "package": "poetry add firecrawl-py", + "packages": ["firecrawl-py"], "env": "FIRECRAWL_API_KEY=...", "tools": ["web_scrape", "web_crawl", "retrieve_web_crawl"], "cta": "Create an API key at https://www.firecrawl.dev/" diff --git a/agentstack/tools/ftp.json b/agentstack/tools/ftp.json index 7a3fd1c..a828de9 100644 --- a/agentstack/tools/ftp.json +++ b/agentstack/tools/ftp.json @@ -1,6 +1,6 @@ { "name": "ftp", - "package": "", + "packages": [], "env": "FTP_HOST=...\nFTP_USER=...\nFTP_PASSWORD=...", "tools": ["upload_files"], "cta": "Be sure to add your FTP credentials to .env" diff --git a/agentstack/tools/mem0.json b/agentstack/tools/mem0.json index 3ccb2b8..a273bc4 100644 --- a/agentstack/tools/mem0.json +++ b/agentstack/tools/mem0.json @@ -1,6 +1,6 @@ { "name": "mem0", - "package": "poetry add mem0ai", + "packages": ["mem0ai"], "env": "MEM0_API_KEY=...", "tools": ["write_to_memory", "read_from_memory"], "cta": "Create your mem0 API key at https://mem0.ai/" diff --git a/agentstack/tools/open_interpreter.json b/agentstack/tools/open_interpreter.json index c709940..a0bdbc8 100644 --- a/agentstack/tools/open_interpreter.json +++ b/agentstack/tools/open_interpreter.json @@ -1,6 +1,6 @@ { "name": "open_interpreter", - "package": "poetry add open-interpreter", + "packages": ["open-interpreter"], "env": "", "tools": ["execute_code"] } \ No newline at end of file diff --git a/agentstack/tools/perplexity.json b/agentstack/tools/perplexity.json index 8e75110..ea35bb3 100644 --- a/agentstack/tools/perplexity.json +++ b/agentstack/tools/perplexity.json @@ -1,6 +1,6 @@ { "name": "perplexity", - "package": "", + "packages": [], "env": "PERPLEXITY_API_KEY=pplx-...", "tools": ["query_perplexity"] } \ No newline at end of file diff --git a/agentstack/tools/vision.json b/agentstack/tools/vision.json index 528dc8e..c622f08 100644 --- a/agentstack/tools/vision.json +++ b/agentstack/tools/vision.json @@ -1,6 +1,6 @@ { "name": "vision", - "package": "poetry add crewai-tools", + "packages": [], "env": "", "tools": ["vision_tool"] } \ No newline at end of file