Clenaup tool_generation, add ToolConfig.module_name, add tests for tool_generation frontend

Author: tcdent

agentstack/generation/tool_generation.py

@@ -1,28 +1,17 @@
 import os, sys
-from typing import Optional, Any, List
-import importlib.resources
+from typing import Optional, Union, Any
 from pathlib import Path
-import json
-import sys
-from typing import Optional, List, Dict, Union
-
-from . import get_agent_names
-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
 import ast
 
+from agentstack import frameworks
 from agentstack import packaging
 from agentstack import ValidationError
-from agentstack.utils import get_package_path
+from agentstack.utils import term_color
 from agentstack.tools import ToolConfig
 from agentstack.generation import astools
 from agentstack.generation.files import ConfigFile, EnvFile
-from agentstack import frameworks
-from .gen_utils import insert_code_after_tag, string_in_file
-from ..utils import open_json_file, get_framework, term_color
 
 
 # This is the filename of the location of tool imports in the user's project.
@@ -44,7 +33,7 @@ def get_import_for_tool(self, tool: ToolConfig) -> ast.Import:
         raises a ValidationError if the tool is imported multiple times.
         """
         all_imports = astools.get_all_imports(self.tree)
-        tool_imports = [i for i in all_imports if tool.name in i.names[0].name]
+        tool_imports = [i for i in all_imports if tool.module_name == i.module]
 
         if len(tool_imports) > 1:
             raise ValidationError(f"Multiple imports for tool {tool.name} found in {self.filename}")
@@ -77,106 +66,100 @@ def remove_import_for_tool(self, framework: str, tool: ToolConfig):
         Remove an import for a tool.
         raises a ValidationError if the tool is not imported.
         """
-        tool_import = self.get_imports_for_tool(tool)
+        tool_import = self.get_import_for_tool(tool)
         if not tool_import:
             raise ValidationError(f"Tool {tool.name} not imported in {self.filename}")
 
         start, end = self.get_node_range(tool_import)
         self.edit_node_range(start, end, "")
 
 
-def add_tool(tool_name: str, agents: Optional[List[str]] = [], path: Optional[str] = None):
-    if path:
-        path = path.endswith('/') and path or path + '/'
-    else:
-        path = './'
-
-    framework = get_framework(path)
+def add_tool(tool_name: str, agents: Optional[list[str]] = [], path: Optional[Path] = None):
+    if path is None: path = Path()
     agentstack_config = ConfigFile(path)
+    framework = agentstack_config.framework
 
     if tool_name in agentstack_config.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 = tool_data.get_impl_file_path(framework)
+    tool = ToolConfig.from_tool_name(tool_name)
+    tool_file_path = tool.get_impl_file_path(framework)
 
-    if tool_data.packages:
-        packaging.install(' '.join(tool_data.packages))
-    shutil.copy(tool_file_path, f'{path}src/tools/{tool_name}_tool.py')  # Move tool from package to project
+    if tool.packages:
+        packaging.install(' '.join(tool.packages))
+    
+    # Move tool from package to project
+    shutil.copy(tool_file_path, path/f'src/tools/{tool.module_name}.py')
 
     try: # Edit the user's project tool init file to include the tool
         with ToolsInitFile(path/TOOLS_INIT_FILENAME) as tools_init:
-            tools_init.add_import_for_tool(tool_data)
+            tools_init.add_import_for_tool(framework, tool)
     except ValidationError as e:
         print(term_color(f"Error adding tool:\n{e}", 'red'))
-        sys.exit(1)
 
     # Edit the framework entrypoint file to include the tool in the agent definition
-    if not len(agents): # If no agents are specified, add the tool to all agents
+    if not agents: # If no agents are specified, add the tool to all agents
         agents = frameworks.get_agent_names(framework, path)
     for agent_name in agents:
-        frameworks.add_tool(framework, tool_data, agent_name, path)
+        frameworks.add_tool(framework, tool, agent_name, path)
 
-    if tool_data.env: # add environment variables which don't exist
+    if tool.env: # add environment variables which don't exist
         with EnvFile(path) as env:
-            for var, value in tool_data.env.items():
+            for var, value in tool.env.items():
                 env.append_if_new(var, value)
         with EnvFile(path, filename=".env.example") as env:
-            for var, value in tool_data.env.items():
+            for var, value in tool.env.items():
                 env.append_if_new(var, value)
 
-    if tool_data.post_install:
-        os.system(tool_data.post_install)
+    if tool.post_install:
+        os.system(tool.post_install)
 
     with agentstack_config as config:
-        config.tools.append(tool_name)
-
-    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'))
+        config.tools.append(tool.name)
 
+    print(term_color(f'🔨 Tool {tool.name} added to agentstack project successfully', 'green'))
+    if tool.cta:
+        print(term_color(f'🪩  {tool.cta}', 'blue'))
 
-def remove_tool(tool_name: str, agents: Optional[List[str]] = [], path: Optional[str] = None):
-    if path:
-        path = path.endswith('/') and path or path + '/'
-    else:
-        path = './'
 
-    framework = get_framework()
+def remove_tool(tool_name: str, agents: Optional[list[str]] = [], path: Optional[Path] = None):
+    if path is None: path = Path()
     agentstack_config = ConfigFile(path)
+    framework = agentstack_config.framework
 
     if not tool_name in agentstack_config.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:
-        packaging.remove(' '.join(tool_data.packages))
+    tool = ToolConfig.from_tool_name(tool_name)
+    if tool.packages:
+        packaging.remove(' '.join(tool.packages))
+
     try:
-        os.remove(f'{path}src/tools/{tool_name}_tool.py')
+        os.remove(path/f'src/tools/{tool.module_name}.py')
     except FileNotFoundError:
-        print(f'"src/tools/{tool_name}_tool.py" not found')
+        print(f'"src/tools/{tool.module_name}.py" not found')
 
     try: # Edit the user's project tool init file to exclude the tool
         with ToolsInitFile(path/TOOLS_INIT_FILENAME) as tools_init:
-            tools_init.remove_import_for_tool(tool_data)
+            tools_init.remove_import_for_tool(framework, tool)
     except ValidationError as e:
         print(term_color(f"Error removing tool:\n{e}", 'red'))
-        sys.exit(1)
 
     # Edit the framework entrypoint file to exclude the tool in the agent definition
-    if not len(agents): # If no agents are specified, remove the tool from all agents
+    if not agents: # If no agents are specified, remove the tool from all agents
         agents = frameworks.get_agent_names(framework, path)
     for agent_name in agents:
-        frameworks.remove_tool(framework, tool_data, agent_name, path)
+        frameworks.remove_tool(framework, tool, agent_name, path)
 
-    if tool_data.post_remove:
-        os.system(tool_data.post_remove)
+    if tool.post_remove:
+        os.system(tool.post_remove)
     # We don't remove the .env variables to preserve user data.
 
     with agentstack_config as config:
-        config.tools.remove(tool_name)
+        config.tools.remove(tool.name)
 
-    print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), term_color('from agentstack project successfully', 'green'))
+    print(term_color(f'🔨 Tool {tool_name}', 'green'), term_color('removed', 'red'), 
+        term_color('from agentstack project successfully', 'green'))
 

agentstack/tools.py

@@ -41,11 +41,15 @@ def from_json(cls, path: Path) -> 'ToolConfig':
                 print(f"{' '.join(error['loc'])}: {error['msg']}")
             sys.exit(1)
 
+    @property
+    def module_name(self) -> str:
+        return f"{self.name}_tool"
+
     def get_import_statement(self, framework: str) -> str:
-        return f"from .{self.name}_tool import {', '.join(self.tools)}"
+        return f"from .{self.module_name} import {', '.join(self.tools)}"
 
     def get_impl_file_path(self, framework: str) -> Path:
-        return get_package_path()/f'templates/{framework}/tools/{self.name}_tool.py'
+        return get_package_path()/f'templates/{framework}/tools/{self.module_name}.py'
 
 def get_all_tool_paths() -> list[Path]:
     paths = []

tests/test_frameworks.py

@@ -3,6 +3,7 @@
 import shutil
 import unittest
 from parameterized import parameterized_class
+
 from agentstack import ValidationError
 from agentstack import frameworks
 from agentstack.tools import ToolConfig

tests/test_tool_generation.py

@@ -0,0 +1,69 @@
+import os, sys
+from pathlib import Path
+import shutil
+import unittest
+from parameterized import parameterized_class
+
+from agentstack import frameworks
+from agentstack.tools import get_all_tools, ToolConfig
+from agentstack.generation.files import ConfigFile
+from agentstack.generation.tool_generation import add_tool, remove_tool, TOOLS_INIT_FILENAME
+
+
+BASE_PATH = Path(__file__).parent
+
+# TODO parameterize all tools
+@parameterized_class([
+    {"framework": framework} for framework in frameworks.SUPPORTED_FRAMEWORKS
+])
+class TestToolGeneration(unittest.TestCase):
+    def setUp(self):
+        self.project_dir = BASE_PATH/'tmp'/'tool_generation'
+        
+        os.makedirs(self.project_dir)
+        os.makedirs(self.project_dir/'src')
+        os.makedirs(self.project_dir/'src'/'tools')
+        (self.project_dir/'src'/'__init__.py').touch()
+        (self.project_dir/TOOLS_INIT_FILENAME).touch()
+        
+        # populate the entrypoint
+        entrypoint_path = frameworks.get_entrypoint_path(self.framework, self.project_dir)
+        shutil.copy(BASE_PATH/f"fixtures/frameworks/{self.framework}/entrypoint_max.py", entrypoint_path)
+        
+        # set the framework in agentstack.json
+        shutil.copy(BASE_PATH/'fixtures'/'agentstack.json', self.project_dir/'agentstack.json')
+        with ConfigFile(self.project_dir) as config:
+            config.framework = self.framework
+    
+    def tearDown(self):
+        shutil.rmtree(self.project_dir)
+    
+    def test_add_tool(self):
+        tool_conf = ToolConfig.from_tool_name('agent_connect')
+        add_tool('agent_connect', path=self.project_dir)
+        
+        entrypoint_path = frameworks.get_entrypoint_path(self.framework, self.project_dir)
+        entrypoint_src = open(entrypoint_path).read()
+        tools_init_src = open(self.project_dir/TOOLS_INIT_FILENAME).read()
+        
+        # TODO verify tool is added to all agents (this is covered in test_frameworks.py)
+        #assert 'agent_connect' in entrypoint_src
+        assert f'from .{tool_conf.module_name} import' in tools_init_src
+        assert (self.project_dir/'src'/'tools'/f'{tool_conf.module_name}.py').exists()
+        assert 'agent_connect' in open(self.project_dir/'agentstack.json').read()
+    
+    def test_remove_tool(self):
+        tool_conf = ToolConfig.from_tool_name('agent_connect')
+        add_tool('agent_connect', path=self.project_dir)
+        remove_tool('agent_connect', path=self.project_dir)
+        
+        entrypoint_path = frameworks.get_entrypoint_path(self.framework, self.project_dir)
+        entrypoint_src = open(entrypoint_path).read()
+        tools_init_src = open(self.project_dir/TOOLS_INIT_FILENAME).read()
+        
+        # TODO verify tool is removed from all agents (this is covered in test_frameworks.py)
+        #assert 'agent_connect' not in entrypoint_src
+        assert f'from .{tool_conf.module_name} import' not in tools_init_src
+        assert not (self.project_dir/'src'/'tools'/f'{tool_conf.module_name}.py').exists()
+        assert 'agent_connect' not in open(self.project_dir/'agentstack.json').read()
+

tests/test_tool_generation_init.py

@@ -3,6 +3,7 @@
 import shutil
 import unittest
 from parameterized import parameterized_class
+
 from agentstack import ValidationError
 from agentstack import frameworks
 from agentstack.tools import ToolConfig
@@ -17,7 +18,7 @@
 ])
 class TestToolGenerationInit(unittest.TestCase):
     def setUp(self):
-        self.project_dir = BASE_PATH/'tmp'/'tool_generation'
+        self.project_dir = BASE_PATH/'tmp'/'tool_generation_init'
         os.makedirs(self.project_dir)
         os.makedirs(self.project_dir/'src')
         os.makedirs(self.project_dir/'src'/'tools')