Merge pull request #49 from tcdent/issue-35

Implement tool removal command #35

Author: bboynton97

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

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
+

agentstack/generation/tool_generation.py

@@ -3,71 +3,139 @@
 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'))
                 if tool_data.get('cta'):
                     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)
 

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:

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"]
 }

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/"

agentstack/tools/code_interpreter.json

@@ -1,6 +1,6 @@
 {
   "name": "code_interpreter",
-  "package": "poetry add crewai-tools",
+  "packages": [],
   "env": "",
   "tools": ["code_interpreter"]
 }

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,

agentstack/tools/directory_search.json

@@ -1,6 +1,6 @@
 {
   "name": "dir_search_tool",
-  "package": "poetry add crewai-tools",
+  "packages": [],
   "env": "",
   "tools": ["dir_search_tool"]
 }

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"