airtrain 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

airtrain/tools/filesystem.py

@@ -0,0 +1,166 @@
+import os
+import json
+import subprocess
+from pathlib import Path
+
+from .registry import StatelessTool, register_tool
+
+
+@register_tool("list_directory")
+class ListDirectoryTool(StatelessTool):
+    """Tool for listing contents of a directory."""
+    
+    def __init__(self):
+        self.name = "list_directory"
+        self.description = "List the contents of a directory, showing files and subdirectories"
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Path to the directory to list. "
+                                  "Defaults to current directory if not provided."
+                },
+                "show_hidden": {
+                    "type": "boolean",
+                    "description": "Whether to show hidden files (starting with .)"
+                }
+            },
+            "required": []
+        }
+        
+    def __call__(self, path: str = ".", show_hidden: bool = False) -> str:
+        """List the contents of a directory."""
+        try:
+            path = os.path.expanduser(path)
+            if not os.path.exists(path):
+                return f"Error: Path '{path}' does not exist"
+            
+            if not os.path.isdir(path):
+                return f"Error: Path '{path}' is not a directory"
+            
+            items = []
+            for item in os.listdir(path):
+                if not show_hidden and item.startswith('.'):
+                    continue
+                    
+                item_path = os.path.join(path, item)
+                item_type = "directory" if os.path.isdir(item_path) else "file"
+                size = os.path.getsize(item_path) if os.path.isfile(item_path) else None
+                
+                items.append({
+                    "name": item,
+                    "type": item_type,
+                    "size": size
+                })
+            
+            return json.dumps({"path": path, "items": items}, indent=2)
+        except Exception as e:
+            return f"Error listing directory: {str(e)}"
+    
+    def to_dict(self):
+        """Convert tool to dictionary format for LLM function calling."""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters
+            }
+        }
+
+
+@register_tool("directory_tree")
+class DirectoryTreeTool(StatelessTool):
+    """Tool for displaying directory structure as a tree."""
+    
+    def __init__(self):
+        self.name = "directory_tree"
+        self.description = "Display the directory structure as a tree, " \
+                          "showing the hierarchy of files and directories"
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Path to the root directory. " \
+                                  "Defaults to current directory if not provided."
+                },
+                "max_depth": {
+                    "type": "integer",
+                    "description": "Maximum depth of subdirectories to display"
+                },
+                "show_hidden": {
+                    "type": "boolean",
+                    "description": "Whether to show hidden files (starting with .)"
+                }
+            },
+            "required": []
+        }
+        
+    def __call__(self, path: str = ".", max_depth: int = 3, show_hidden: bool = False) -> str:
+        """Display the directory structure as a tree."""
+        try:
+            path = os.path.expanduser(path)
+            if not os.path.exists(path):
+                return f"Error: Path '{path}' does not exist"
+            
+            if not os.path.isdir(path):
+                return f"Error: Path '{path}' is not a directory"
+            
+            # Try to use external 'tree' command if available
+            try:
+                cmd = ["tree", path, "-L", str(max_depth)]
+                if not show_hidden:
+                    cmd.append("-I")
+                    cmd.append(".*")
+                
+                result = subprocess.run(cmd, capture_output=True, text=True)
+                if result.returncode == 0:
+                    return result.stdout
+            except (subprocess.SubprocessError, FileNotFoundError):
+                # If 'tree' command fails or is not available, fall back to custom implementation
+                pass
+                
+            # Custom tree implementation
+            result = [f"Directory tree for {path}:"]
+            
+            def add_to_tree(directory: Path, prefix: str = "", depth: int = 0):
+                if depth > max_depth:
+                    return
+                
+                try:
+                    entries = sorted(directory.iterdir(), 
+                                   key=lambda x: (x.is_file(), x.name))
+                    
+                    for i, entry in enumerate(entries):
+                        if not show_hidden and entry.name.startswith('.'):
+                            continue
+                            
+                        is_last = i == len(entries) - 1
+                        result.append(f"{prefix}{'└── ' if is_last else '├── '}{entry.name}")
+                        
+                        if entry.is_dir():
+                            add_to_tree(
+                                entry,
+                                prefix + ('    ' if is_last else '│   '),
+                                depth + 1
+                            )
+                except PermissionError:
+                    result.append(f"{prefix}└── [Permission denied]")
+            
+            add_to_tree(Path(path))
+            return "\n".join(result)
+        except Exception as e:
+            return f"Error creating directory tree: {str(e)}"
+    
+    def to_dict(self):
+        """Convert tool to dictionary format for LLM function calling."""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters
+            }
+        } 

airtrain/tools/network.py

@@ -0,0 +1,111 @@
+import requests
+from typing import Dict, Any, Optional, Union
+from urllib.parse import urlparse
+
+from .registry import StatelessTool, register_tool
+
+
+@register_tool("api_call")
+class ApiCallTool(StatelessTool):
+    """Tool for making HTTP API calls."""
+    
+    def __init__(self):
+        self.name = "api_call"
+        self.description = "Make HTTP requests to external APIs"
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "url": {
+                    "type": "string",
+                    "description": "URL to make the request to"
+                },
+                "method": {
+                    "type": "string",
+                    "enum": ["GET", "POST", "PUT", "DELETE", "PATCH"],
+                    "description": "HTTP method to use for the request"
+                },
+                "headers": {
+                    "type": "object",
+                    "description": "HTTP headers to include in the request"
+                },
+                "params": {
+                    "type": "object",
+                    "description": "URL parameters for the request"
+                },
+                "data": {
+                    "type": "object",
+                    "description": "Data to send in the request body (for POST, PUT, PATCH)"
+                },
+                "json_data": {
+                    "type": "object",
+                    "description": "JSON data to send in the request body (for POST, PUT, PATCH)"
+                },
+                "timeout": {
+                    "type": "number",
+                    "description": "Request timeout in seconds"
+                }
+            },
+            "required": ["url", "method"]
+        }
+        
+    def __call__(
+        self, 
+        url: str, 
+        method: str = "GET", 
+        headers: Optional[Dict[str, str]] = None,
+        params: Optional[Dict[str, str]] = None,
+        data: Optional[Union[Dict[str, Any], str]] = None,
+        json_data: Optional[Dict[str, Any]] = None,
+        timeout: float = 10.0
+    ) -> Dict[str, Any]:
+        """Make an HTTP request to the specified URL."""
+        try:
+            # Validate URL
+            parsed_url = urlparse(url)
+            if not parsed_url.scheme or not parsed_url.netloc:
+                return {"error": f"Invalid URL '{url}'"}
+            
+            # Prepare request
+            method = method.upper()
+            if method not in ["GET", "POST", "PUT", "DELETE", "PATCH"]:
+                return {"error": f"Unsupported HTTP method '{method}'"}
+            
+            # Make request
+            response = requests.request(
+                method=method,
+                url=url,
+                headers=headers,
+                params=params,
+                data=data,
+                json=json_data,
+                timeout=timeout
+            )
+            
+            # Try to parse response as JSON
+            try:
+                json_result = response.json()
+                response_data = json_result
+            except ValueError:
+                # Not JSON, return text
+                response_data = response.text
+            
+            return {
+                "status_code": response.status_code,
+                "headers": dict(response.headers),
+                "content": response_data
+            }
+        except requests.exceptions.RequestException as e:
+            return {"error": f"Error making API request: {str(e)}"}
+        except Exception as e:
+            return {"error": f"Error: {str(e)}"}
+    
+    def to_dict(self):
+        """Convert tool to dictionary format for LLM function calling."""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters
+            }
+        } 

airtrain/tools/registry.py

@@ -0,0 +1,320 @@
+"""
+Tool Registry System for AirTrain
+
+This module provides a registry system for tools that can be used by AI agents.
+It supports both stateful tools (requiring fresh instances) and stateless tools 
+(which can be shared/reused).
+
+The registry system includes:
+- Validation mechanisms for tools
+- Registration decorators
+- Factory methods for tool creation
+- Discovery utilities for finding available tools
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List, Type, Any, Optional, TypeVar
+
+
+# Type variable for tool classes
+T = TypeVar('T', bound='BaseTool')
+
+# Registry structure: {"stateful": {}, "stateless": {}}
+TOOL_REGISTRY = {
+    "stateful": {},
+    "stateless": {}
+}
+
+
+class ToolValidationError(Exception):
+    """Exception raised when a tool fails validation checks."""
+    pass
+
+
+class BaseTool(ABC):
+    """Base class for all tools."""
+    
+    # These will be set by the registration decorator
+    tool_name: str = None
+    tool_type: str = None
+    
+    @abstractmethod
+    def __call__(self, **kwargs) -> Any:
+        """Execute the tool with the given parameters."""
+        pass
+    
+    @abstractmethod
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert tool to dictionary format for LLM function calling."""
+        pass
+
+
+class StatefulTool(BaseTool):
+    """Base class for tools that maintain state and require fresh instances."""
+    
+    @classmethod
+    @abstractmethod
+    def create_instance(cls: Type[T]) -> T:
+        """Create a new instance of the tool with fresh state."""
+        pass
+    
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the tool's internal state."""
+        pass
+
+
+class StatelessTool(BaseTool):
+    """Base class for stateless tools that can be reused."""
+    pass
+
+
+def validate_tool(cls: Type[BaseTool], tool_type: str) -> Type[BaseTool]:
+    """
+    Validate that a tool class meets the requirements for its type.
+    
+    Args:
+        cls: The tool class to validate
+        tool_type: Either "stateful" or "stateless"
+    
+    Returns:
+        The validated tool class
+    
+    Raises:
+        ToolValidationError: If the tool does not meet requirements
+    """
+    # Check that the class implements the required methods
+    if not hasattr(cls, '__call__') or not callable(getattr(cls, '__call__')):
+        raise ToolValidationError(f"Tool {cls.__name__} must implement __call__ method")
+    
+    if not hasattr(cls, 'to_dict') or not callable(getattr(cls, 'to_dict')):
+        raise ToolValidationError(f"Tool {cls.__name__} must implement to_dict method")
+    
+    # Validate stateful tool specific requirements
+    if tool_type == "stateful":
+        if not issubclass(cls, StatefulTool):
+            raise ToolValidationError(
+                f"Stateful tool {cls.__name__} must inherit from StatefulTool"
+            )
+        
+        create_instance_attr = hasattr(cls, 'create_instance')
+        create_instance_callable = callable(getattr(cls, 'create_instance', None))
+        
+        if not create_instance_attr or not create_instance_callable:
+            raise ToolValidationError(
+                f"Stateful tool {cls.__name__} must implement create_instance class method"
+            )
+        
+        if not hasattr(cls, 'reset') or not callable(getattr(cls, 'reset')):
+            raise ToolValidationError(
+                f"Stateful tool {cls.__name__} must implement reset method"
+            )
+    
+    # Validate stateless tool specific requirements
+    if tool_type == "stateless":
+        if not issubclass(cls, StatelessTool):
+            raise ToolValidationError(
+                f"Stateless tool {cls.__name__} must inherit from StatelessTool"
+            )
+    
+    return cls
+
+
+def register_tool(name: str, tool_type: str = "stateless"):
+    """
+    Decorator for registering a tool with the registry.
+    
+    Args:
+        name: The name of the tool
+        tool_type: Either "stateful" or "stateless"
+    
+    Returns:
+        A decorator function that registers the tool
+    
+    Raises:
+        ValueError: If the tool name is already registered or the tool type is invalid
+    """
+    if tool_type not in TOOL_REGISTRY:
+        raise ValueError(
+            f"Invalid tool type: {tool_type}. Must be either 'stateful' or 'stateless'"
+        )
+    
+    def decorator(cls: Type[BaseTool]) -> Type[BaseTool]:
+        if name in TOOL_REGISTRY[tool_type]:
+            raise ValueError(f"Tool name '{name}' already registered in {tool_type} registry")
+        
+        # Validate the tool
+        validated_cls = validate_tool(cls, tool_type)
+        
+        # Register the tool
+        TOOL_REGISTRY[tool_type][name] = validated_cls
+        
+        # Add metadata to the class
+        validated_cls.tool_name = name
+        validated_cls.tool_type = tool_type
+        
+        return validated_cls
+    
+    return decorator
+
+
+class ToolFactory:
+    """Factory class for creating and managing tools."""
+    
+    @staticmethod
+    def get_tool(name: str, tool_type: str = "stateless") -> BaseTool:
+        """
+        Get a tool instance by name and type.
+        
+        For stateful tools, this returns a fresh instance.
+        For stateless tools, this returns a singleton instance.
+        
+        Args:
+            name: The name of the tool
+            tool_type: Either "stateful" or "stateless"
+        
+        Returns:
+            An instance of the requested tool
+        
+        Raises:
+            ValueError: If the tool or tool type is not found
+        """
+        if tool_type not in TOOL_REGISTRY:
+            raise ValueError(f"Invalid tool type: {tool_type}")
+        
+        tool_cls = TOOL_REGISTRY[tool_type].get(name)
+        if not tool_cls:
+            raise ValueError(f"Tool '{name}' not found in {tool_type} registry")
+        
+        # Handle stateful tools - always create a fresh instance
+        if tool_type == "stateful":
+            instance = tool_cls.create_instance()
+            instance.reset()  # Ensure the instance is in a clean state
+            return instance
+        
+        # Handle stateless tools - reuse the same instance
+        # We use a singleton pattern here with lazy initialization
+        if not hasattr(tool_cls, '_instance'):
+            tool_cls._instance = tool_cls()
+        
+        return tool_cls._instance
+    
+    @staticmethod
+    def list_tools(tool_type: Optional[str] = None) -> Dict[str, List[str]]:
+        """
+        List all registered tools, optionally filtered by type.
+        
+        Args:
+            tool_type: Optional filter for tool type
+        
+        Returns:
+            A dictionary mapping tool types to lists of tool names
+        """
+        if tool_type:
+            if tool_type not in TOOL_REGISTRY:
+                raise ValueError(f"Invalid tool type: {tool_type}")
+            return {tool_type: list(TOOL_REGISTRY[tool_type].keys())}
+        
+        return {t_type: list(tools.keys()) for t_type, tools in TOOL_REGISTRY.items()}
+    
+    @staticmethod
+    def get_tool_definitions(tool_type: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        Get tool definitions for all registered tools or tools of a specific type.
+        
+        This is useful for preparing tools for LLM function calling.
+        
+        Args:
+            tool_type: Optional filter for tool type
+        
+        Returns:
+            A list of tool definitions in dictionary format
+        """
+        tool_defs = []
+        
+        if tool_type:
+            if tool_type not in TOOL_REGISTRY:
+                raise ValueError(f"Invalid tool type: {tool_type}")
+            registry = {tool_type: TOOL_REGISTRY[tool_type]}
+        else:
+            registry = TOOL_REGISTRY
+        
+        for t_type, tools in registry.items():
+            for name, cls in tools.items():
+                # For stateless tools, we can use the singleton instance
+                if t_type == "stateless":
+                    if not hasattr(cls, '_instance'):
+                        cls._instance = cls()
+                    tool_defs.append(cls._instance.to_dict())
+                # For stateful tools, we need to create a temporary instance
+                else:
+                    instance = cls.create_instance()
+                    tool_defs.append(instance.to_dict())
+        
+        return tool_defs
+
+
+def get_default_tools(tool_type: Optional[str] = None) -> List[Dict[str, Any]]:
+    """
+    Get tool definitions for all registered tools.
+    
+    This is a convenience function that delegates to ToolFactory.get_tool_definitions.
+    
+    Args:
+        tool_type: Optional filter for tool type
+    
+    Returns:
+        A list of tool definitions in dictionary format
+    """
+    return ToolFactory.get_tool_definitions(tool_type)
+
+
+def execute_tool_call(tool_call: Dict[str, Any]) -> Any:
+    """
+    Execute a tool call based on LLM function calling format.
+    
+    Args:
+        tool_call: A dictionary containing the tool call details
+    
+    Returns:
+        The result of executing the tool call
+    
+    Raises:
+        ValueError: If the tool is not found or the call format is invalid
+    """
+    import json
+    
+    # Extract tool details from the call
+    function_details = tool_call.get("function", {})
+    function_name = function_details.get("name")
+    
+    if not function_name:
+        raise ValueError("Invalid tool call format: Missing function name")
+    
+    # Try to find the tool in both registries
+    tool = None
+    tool_type = None
+    
+    for t_type in TOOL_REGISTRY:
+        if function_name in TOOL_REGISTRY[t_type]:
+            tool_type = t_type
+            break
+    
+    if not tool_type:
+        raise ValueError(f"Tool '{function_name}' not found in any registry")
+    
+    # Get a tool instance
+    tool = ToolFactory.get_tool(function_name, tool_type)
+    
+    # Parse arguments
+    try:
+        arguments = json.loads(function_details.get("arguments", "{}"))
+    except json.JSONDecodeError:
+        raise ValueError(f"Invalid arguments format for tool '{function_name}'")
+    
+    # Execute the tool
+    try:
+        result = tool(**arguments)
+        return result
+    except Exception as e:
+        return f"Error executing tool '{function_name}': {str(e)}"