created better throw exception logic for tools / tool_generation.py

agentstack/cli/tools.py

@@ -1,31 +1,43 @@
 from typing import Optional
 import itertools
+import sys
 import inquirer
 from agentstack.utils import term_color
 from agentstack import generation
 from agentstack.tools import get_all_tools
 from agentstack.agents import get_all_agents
+from agentstack.exceptions import ToolError, ValidationError
 
 
 def list_tools():
     """
     List all available tools by category.
     """
-    tools = get_all_tools()
-    curr_category = None
+    try:
+        tools = get_all_tools()
+        curr_category = None
 
-    print("\n\nAvailable AgentStack Tools:")
-    for category, tools in itertools.groupby(tools, lambda x: x.category):
-        if curr_category != category:
-            print(f"\n{category}:")
-            curr_category = category
-        for tool in tools:
-            print("  - ", end='')
-            print(term_color(f"{tool.name}", 'blue'), end='')
-            print(f": {tool.url if tool.url else 'AgentStack default tool'}")
+        print("\n\nAvailable AgentStack Tools:")
+        for category, tools in itertools.groupby(tools, lambda x: x.category):
+            if curr_category != category:
+                print(f"\n{category}:")
+                curr_category = category
+            for tool in tools:
+                print("  - ", end='')
+                print(term_color(f"{tool.name}", 'blue'), end='')
+                print(f": {tool.url if tool.url else 'AgentStack default tool'}")
 
-    print("\n\n✨ Add a tool with: agentstack tools add <tool_name>")
-    print("   https://docs.agentstack.sh/tools/core")
+        print("\n\n✨ Add a tool with: agentstack tools add <tool_name>")
+        print("   https://docs.agentstack.sh/tools/core")
+    except ToolError:
+        print(term_color("Could not retrieve list of tools. The tools directory may be corrupted or missing.", 'red'))
+        sys.exit(1)
+    except ValidationError as e:
+        print(term_color(f"Validation error: {str(e)}", 'red'))
+        sys.exit(1)
+    except Exception:
+        print(term_color("An unexpected error occurred while listing tools.", 'red'))
+        sys.exit(1)
 
 
 def add_tool(tool_name: Optional[str], agents=Optional[list[str]]):
@@ -38,32 +50,60 @@ def add_tool(tool_name: Optional[str], agents=Optional[list[str]]):
         - add the tool to the user's project
         - add the tool to the specified agents or all agents if none are specified
     """
-    if not tool_name:
-        # ask the user for the tool name
-        tools_list = [
-            inquirer.List(
-                "tool_name",
-                message="Select a tool to add to your project",
-                choices=[tool.name for tool in get_all_tools()],
-            )
-        ]
-        try:
-            tool_name = inquirer.prompt(tools_list)['tool_name']
-        except TypeError:
-            return  # user cancelled the prompt
+    try:
+        if not tool_name:
+            # ask the user for the tool name
+            tools_list = [
+                inquirer.List(
+                    "tool_name",
+                    message="Select a tool to add to your project",
+                    choices=[tool.name for tool in get_all_tools()],
+                )
+            ]
+            try:
+                tool_name = inquirer.prompt(tools_list)['tool_name']
+            except TypeError:
+                return  # user cancelled the prompt
 
-        # ask the user for the agents to add the tool to
-        agents_list = [
-            inquirer.Checkbox(
-                "agents",
-                message="Select which agents to make the tool available to",
-                choices=[agent.name for agent in get_all_agents()],
-            )
-        ]
-        try:
-            agents = inquirer.prompt(agents_list)['agents']
-        except TypeError:
-            return  # user cancelled the prompt
+            # ask the user for the agents to add the tool to
+            agents_list = [
+                inquirer.Checkbox(
+                    "agents",
+                    message="Select which agents to make the tool available to",
+                    choices=[agent.name for agent in get_all_agents()],
+                )
+            ]
+            try:
+                agents = inquirer.prompt(agents_list)['agents']
+            except TypeError:
+                return  # user cancelled the prompt
 
-    assert tool_name  # appease type checker
-    generation.add_tool(tool_name, agents=agents)
+        assert tool_name  # appease type checker
+        generation.add_tool(tool_name, agents=agents)
+    except ToolError:
+        print(term_color(f"Could not add tool '{tool_name}'. Run 'agentstack tools list' to see available tools.", 'red'))
+        sys.exit(1)
+    except ValidationError as e:
+        print(term_color(f"Validation error: {str(e)}", 'red'))
+        sys.exit(1)
+    except Exception:
+        print(term_color("An unexpected error occurred while adding the tool.", 'red'))
+        sys.exit(1)
+
+
+def remove_tool(tool_name: str, agents: Optional[list[str]] = []):
+    """Remove a tool from the project"""
+    try:
+        generation.remove_tool(tool_name, agents=agents)
+    except ToolError as e:
+        if "not installed" in str(e):
+            print(term_color(f"Tool '{tool_name}' is not installed in this project.", 'red'))
+        else:
+            print(term_color(f"Could not remove tool '{tool_name}'. The tool may be in use or corrupted.", 'red'))
+        sys.exit(1)
+    except ValidationError as e:
+        print(term_color(f"Validation error: {str(e)}", 'red'))
+        sys.exit(1)
+    except Exception:
+        print(term_color("An unexpected error occurred while removing the tool.", 'red'))
+        sys.exit(1)

agentstack/exceptions.py

@@ -1,7 +1,14 @@
+from pathlib import Path
 class ValidationError(Exception):
     """
     Raised when a validation error occurs ie. a file does not meet the required
     format or a syntax error is found.
     """
 
     pass
+
+
+class ToolError(Exception):
+    """Base exception for tool-related errors"""
+
+    pass

agentstack/generation/tool_generation.py

@@ -1,13 +1,12 @@
 import os
-import sys
 from typing import Optional
 from pathlib import Path
 import shutil
 import ast
 
 from agentstack import conf
 from agentstack.conf import ConfigFile
-from agentstack.exceptions import ValidationError
+from agentstack.exceptions import ValidationError, ToolError
 from agentstack import frameworks
 from agentstack import packaging
 from agentstack.utils import term_color
@@ -97,7 +96,7 @@ def add_tool(tool_name: str, agents: Optional[list[str]] = []):
             with ToolsInitFile(conf.PATH / TOOLS_INIT_FILENAME) as tools_init:
                 tools_init.add_import_for_tool(tool, agentstack_config.framework)
         except ValidationError as e:
-            print(term_color(f"Error adding tool:\n{e}", 'red'))
+            raise  # Let ValidationError propagate up to CLI layer for proper error handling
 
         if tool.env:  # add environment variables which don't exist
             with EnvFile() as env:
@@ -129,8 +128,7 @@ def remove_tool(tool_name: str, agents: Optional[list[str]] = []):
     agentstack_config = ConfigFile()
 
     if tool_name not in agentstack_config.tools:
-        print(term_color(f'Tool {tool_name} is not installed', 'red'))
-        sys.exit(1)
+        raise ToolError(f'Tool {tool_name} is not installed')
 
     tool = ToolConfig.from_tool_name(tool_name)
     if tool.packages:
@@ -146,7 +144,7 @@ def remove_tool(tool_name: str, agents: Optional[list[str]] = []):
         with ToolsInitFile(conf.PATH / TOOLS_INIT_FILENAME) as tools_init:
             tools_init.remove_import_for_tool(tool, agentstack_config.framework)
     except ValidationError as e:
-        print(term_color(f"Error removing tool:\n{e}", 'red'))
+        raise  # Let ValidationError propagate up to CLI layer for proper error handling
 
     # Edit the framework entrypoint file to exclude the tool in the agent definition
     if not agents:  # If no agents are specified, remove the tool from all agents

agentstack/tools.py

@@ -1,9 +1,9 @@
 from typing import Optional
 import os
-import sys
 from pathlib import Path
 import pydantic
 from agentstack.utils import get_package_path, open_json_file, term_color
+from agentstack.exceptions import ValidationError, ToolError
 
 
 class ToolConfig(pydantic.BaseModel):
@@ -26,9 +26,8 @@ class ToolConfig(pydantic.BaseModel):
     @classmethod
     def from_tool_name(cls, name: str) -> 'ToolConfig':
         path = get_package_path() / f'tools/{name}.json'
-        if not os.path.exists(path):  # TODO raise exceptions and handle message/exit in cli
-            print(term_color(f'No known agentstack tool: {name}', 'red'))
-            sys.exit(1)
+        if not os.path.exists(path):
+            raise ToolError(f'No known agentstack tool: {name}')
         return cls.from_json(path)
 
     @classmethod
@@ -37,11 +36,9 @@ def from_json(cls, path: Path) -> 'ToolConfig':
         try:
             return cls(**data)
         except pydantic.ValidationError as e:
-            # TODO raise exceptions and handle message/exit in cli
-            print(term_color(f"Error validating tool config JSON: \n{path}", 'red'))
-            for error in e.errors():
-                print(f"{' '.join([str(loc) for loc in error['loc']])}: {error['msg']}")
-            sys.exit(1)
+            error_msg = f"Error validating tool config JSON: \n{path}\n"
+            error_msg += "\n".join(f"{' '.join([str(loc) for loc in error['loc']])}: {error['msg']}" for error in e.errors())
+            raise ValidationError(error_msg)
 
     @property
     def module_name(self) -> str: