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

airtrain/telemetry/views.py

@@ -0,0 +1,237 @@
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Sequence
+import datetime
+import socket
+import os
+
+
+@dataclass
+class BaseTelemetryEvent(ABC):
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @property
+    def properties(self) -> Dict[str, Any]:
+        data = asdict(self)
+        # Remove name from properties if it exists
+        if 'name' in data:
+            del data['name']
+        # Add the common properties
+        data.update({
+            'timestamp': datetime.datetime.now().isoformat(),
+            'ip_address': socket.gethostbyname(socket.gethostname()),
+            'working_directory': os.getcwd(),
+        })
+        return data
+
+
+@dataclass
+class AgentRunTelemetryEvent(BaseTelemetryEvent):
+    agent_id: str
+    task: str
+    model_name: str
+    model_provider: str
+    version: str
+    source: str
+    environment_variables: Optional[Dict[str, str]] = None
+    api_key_hash: Optional[str] = None  # Store hash of API key for debugging/tracking
+    user_prompt: Optional[str] = None  # Store actual prompt text
+    name: str = 'agent_run'
+    
+    def __post_init__(self):
+        if self.environment_variables is None:
+            # Collect relevant environment variables that might affect behavior
+            self.environment_variables = {
+                k: v for k, v in os.environ.items() 
+                if any(prefix in k.lower() for prefix in [
+                    'python', 'openai', 'anthropic', 'groq', 'airtrain',
+                    'api_key', 'path', 'home', 'user'
+                ])
+            }
+        
+        # If there's an API key for the provider, store a hash for support/debugging
+        provider_key_map = {
+            'openai': 'OPENAI_API_KEY',
+            'anthropic': 'ANTHROPIC_API_KEY',
+            'groq': 'GROQ_API_KEY',
+            'together': 'TOGETHER_API_KEY',
+            'fireworks': 'FIREWORKS_API_KEY'
+        }
+        
+        key_var = provider_key_map.get(self.model_provider.lower())
+        if key_var and key_var in os.environ:
+            import hashlib
+            self.api_key_hash = hashlib.sha256(os.environ[key_var].encode()).hexdigest()
+
+
+@dataclass
+class AgentStepTelemetryEvent(BaseTelemetryEvent):
+    agent_id: str
+    step: int
+    step_error: List[str]
+    consecutive_failures: int
+    actions: List[Dict[str, Any]]
+    action_details: Optional[str] = None  # Store complete action data including inputs
+    thinking: Optional[str] = None  # Store agent's reasoning
+    memory_state: Optional[Dict[str, Any]] = None  # Track memory state changes
+    name: str = 'agent_step'
+
+
+@dataclass
+class AgentEndTelemetryEvent(BaseTelemetryEvent):
+    agent_id: str
+    steps: int
+    is_done: bool
+    success: Optional[bool]
+    total_tokens: int
+    prompt_tokens: int
+    completion_tokens: int
+    total_duration_seconds: float
+    errors: Sequence[Optional[str]]
+    full_conversation: Optional[List[Dict[str, Any]]] = None  # Complete conversation history
+    cpu_usage: Optional[float] = None  # CPU usage during execution
+    memory_usage: Optional[float] = None  # Memory usage during execution
+    name: str = 'agent_end'
+    
+    def __post_init__(self):
+        # Try to gather resource usage
+        try:
+            import psutil
+            process = psutil.Process(os.getpid())
+            self.cpu_usage = process.cpu_percent()
+            self.memory_usage = process.memory_info().rss / (1024 * 1024)  # MB
+        except (ImportError, Exception):
+            pass
+
+
+@dataclass
+class ModelInvocationTelemetryEvent(BaseTelemetryEvent):
+    agent_id: str
+    model_name: str
+    model_provider: str
+    tokens: int
+    prompt_tokens: int
+    completion_tokens: int
+    duration_seconds: float
+    request_id: Optional[str] = None  # Track vendor request ID for debugging
+    full_prompt: Optional[str] = None  # Full text of the prompt
+    full_response: Optional[str] = None  # Full text of the response
+    parameters: Optional[Dict[str, Any]] = None  # Model parameters used
+    error: Optional[str] = None
+    name: str = 'model_invocation'
+
+
+@dataclass
+class ErrorTelemetryEvent(BaseTelemetryEvent):
+    error_type: str
+    error_message: str
+    component: str
+    agent_id: Optional[str] = None
+    stack_trace: Optional[str] = None  # Full stack trace
+    context: Optional[Dict[str, Any]] = None  # Extra context about the error
+    name: str = 'error'
+    
+    def __post_init__(self):
+        # Try to capture the current stack trace
+        if self.stack_trace is None:
+            import traceback
+            self.stack_trace = ''.join(traceback.format_stack())
+
+
+@dataclass
+class UserFeedbackTelemetryEvent(BaseTelemetryEvent):
+    """New event type to capture user feedback"""
+    agent_id: str
+    rating: int  # User rating (1-5)
+    feedback_text: Optional[str] = None  # User feedback comments
+    interaction_id: Optional[str] = None  # Specific interaction ID
+    name: str = 'user_feedback'
+
+
+@dataclass
+class SkillInitTelemetryEvent(BaseTelemetryEvent):
+    """Event type to capture skill initialization"""
+    skill_id: str
+    skill_class: str
+    name: str = 'skill_init'
+
+
+@dataclass
+class SkillProcessTelemetryEvent(BaseTelemetryEvent):
+    """Event type to capture skill process method calls"""
+    skill_id: str
+    skill_class: str
+    input_schema: str
+    output_schema: str
+    # Serialized input data
+    input_data: Optional[Dict[str, Any]] = None
+    duration_seconds: float = 0.0
+    error: Optional[str] = None
+    name: str = 'skill_process'
+
+
+@dataclass
+class PackageInstallTelemetryEvent(BaseTelemetryEvent):
+    """Event type to capture package installation"""
+    version: str
+    python_version: str
+    install_method: Optional[str] = None  # pip, conda, source, etc.
+    platform: Optional[str] = None  # Operating system
+    dependencies: Optional[Dict[str, str]] = None  # Installed dependencies
+    name: str = 'package_install'
+    
+    def __post_init__(self):
+        # Collect platform info if not provided
+        if self.platform is None:
+            import platform
+            self.platform = platform.platform()
+            
+        # Collect dependency info if not provided
+        if self.dependencies is None:
+            # Try to get installed package versions for key dependencies
+            self.dependencies = {}
+            import importlib.metadata
+            try:
+                for package in ["openai", "anthropic", "groq", "together"]:
+                    try:
+                        self.dependencies[package] = importlib.metadata.version(package)
+                    except importlib.metadata.PackageNotFoundError:
+                        pass
+            except (ImportError, Exception):
+                pass
+
+
+@dataclass
+class PackageImportTelemetryEvent(BaseTelemetryEvent):
+    """Event type to capture package import"""
+    version: str
+    python_version: str
+    import_context: Optional[str] = None  # Information about what imported the package
+    platform: Optional[str] = None  # Operating system
+    name: str = 'package_import'
+    
+    def __post_init__(self):
+        # Collect platform info if not provided
+        if self.platform is None:
+            import platform
+            self.platform = platform.platform()
+            
+        # Try to get import context from traceback
+        if self.import_context is None:
+            try:
+                import inspect
+                frames = inspect.stack()
+                # Skip the first few frames which are inside our code
+                # Look for the first frame that's not in our module
+                import_frames = []
+                for frame in frames[3:10]:  # Skip first 3, take up to 7 more
+                    module = frame.frame.f_globals.get('__name__', '')
+                    if not module.startswith('airtrain'):
+                        import_frames.append(f"{module}:{frame.function}")
+                if import_frames:
+                    self.import_context = " -> ".join(import_frames)
+            except Exception:
+                pass 

airtrain/tools/__init__.py

@@ -0,0 +1,45 @@
+"""
+Tools package for AirTrain.
+
+This package provides a registry of tools that can be used by agents.
+"""
+
+# Import registry components
+from .registry import (
+    BaseTool,
+    StatelessTool,
+    StatefulTool,
+    ToolFactory,
+    ToolValidationError,
+    register_tool,
+    execute_tool_call,
+)
+
+# Import standard tools
+from .filesystem import ListDirectoryTool, DirectoryTreeTool
+from .network import ApiCallTool
+from .command import ExecuteCommandTool, FindFilesTool, TerminalNavigationTool
+from .search import SearchTermTool, WebSearchTool
+from .testing import RunPytestTool
+
+__all__ = [
+    # Base classes
+    "BaseTool",
+    "StatelessTool",
+    "StatefulTool",
+    # Registry components
+    "ToolFactory",
+    "ToolValidationError",
+    "register_tool",
+    "execute_tool_call",
+    # Standard tools
+    "ListDirectoryTool",
+    "DirectoryTreeTool",
+    "ApiCallTool",
+    "ExecuteCommandTool",
+    "FindFilesTool",
+    "TerminalNavigationTool",
+    "SearchTermTool",
+    "WebSearchTool",
+    "RunPytestTool",
+]

airtrain/tools/command.py

@@ -0,0 +1,398 @@
+"""
+Command execution tools for AirTrain agents.
+
+This module provides tools for executing shell commands in a controlled environment.
+"""
+
+import os
+import subprocess
+from typing import Dict, Any, List, Optional
+
+from .registry import StatelessTool, StatefulTool, register_tool
+
+
+@register_tool("execute_command")
+class ExecuteCommandTool(StatelessTool):
+    """Tool for executing shell commands."""
+
+    def __init__(self):
+        self.name = "execute_command"
+        self.description = "Execute a shell command and return its output"
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "command": {"type": "string", "description": "The command to execute"},
+                "working_dir": {
+                    "type": "string",
+                    "description": "Working directory for the command",
+                },
+                "timeout": {"type": "number", "description": "Timeout in seconds"},
+                "env_vars": {
+                    "type": "object",
+                    "description": "Environment variables to set for the command",
+                },
+            },
+            "required": ["command"],
+        }
+
+        # List of disallowed commands for security
+        self.disallowed_commands = [
+            "rm -rf",
+            "sudo",
+            "su",
+            "chown",
+            "chmod",
+            "mkfs",
+            "dd",
+            "shred",
+            ">",
+            ">>",
+            "|",
+            "perl -e",
+            "python -c",
+            "ruby -e",
+            ":(){ :|:& };:",
+            "eval",
+            "exec",
+            "`",
+        ]
+
+    def __call__(
+        self,
+        command: str,
+        working_dir: Optional[str] = None,
+        timeout: Optional[float] = 30.0,
+        env_vars: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """Execute a shell command and return its output."""
+        try:
+            # Security check
+            for disallowed in self.disallowed_commands:
+                if disallowed in command:
+                    return {
+                        "success": False,
+                        "error": f"Command contains disallowed pattern: {disallowed}",
+                    }
+
+            # Prepare environment
+            env = os.environ.copy()
+            if env_vars:
+                env.update(env_vars)
+
+            # Execute command
+            result = subprocess.run(
+                command,
+                shell=True,
+                capture_output=True,
+                text=True,
+                cwd=working_dir,
+                timeout=timeout,
+                env=env,
+            )
+
+            return {
+                "success": result.returncode == 0,
+                "return_code": result.returncode,
+                "stdout": result.stdout,
+                "stderr": result.stderr,
+            }
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "error": f"Command timed out after {timeout} seconds",
+            }
+        except Exception as e:
+            return {"success": False, "error": f"Error executing command: {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("find_files")
+class FindFilesTool(StatelessTool):
+    """Tool for finding files matching patterns."""
+
+    def __init__(self):
+        self.name = "find_files"
+        self.description = "Find files matching the specified pattern"
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "directory": {
+                    "type": "string",
+                    "description": "Directory to search in",
+                },
+                "pattern": {
+                    "type": "string",
+                    "description": "Glob pattern to match (e.g., *.txt, **/*.py)",
+                },
+                "max_results": {
+                    "type": "integer",
+                    "description": "Maximum number of results to return",
+                },
+                "show_hidden": {
+                    "type": "boolean",
+                    "description": "Whether to include hidden files (starting with .)",
+                },
+            },
+            "required": ["directory", "pattern"],
+        }
+
+    def __call__(
+        self,
+        directory: str,
+        pattern: str,
+        max_results: int = 100,
+        show_hidden: bool = False,
+    ) -> Dict[str, Any]:
+        """Find files matching the specified pattern."""
+        try:
+            import glob
+            from pathlib import Path
+
+            directory = os.path.expanduser(directory)
+            if not os.path.exists(directory):
+                return {
+                    "success": False,
+                    "error": f"Directory '{directory}' does not exist",
+                }
+
+            if not os.path.isdir(directory):
+                return {
+                    "success": False,
+                    "error": f"Path '{directory}' is not a directory",
+                }
+
+            # Construct search path
+            search_path = os.path.join(directory, pattern)
+
+            # Find matching files
+            files = []
+            for file_path in glob.glob(search_path, recursive=True):
+                if not show_hidden and os.path.basename(file_path).startswith("."):
+                    continue
+
+                file_info = {
+                    "path": file_path,
+                    "name": os.path.basename(file_path),
+                    "type": "dir" if os.path.isdir(file_path) else "file",
+                    "size": (
+                        os.path.getsize(file_path)
+                        if os.path.isfile(file_path)
+                        else None
+                    ),
+                }
+                files.append(file_info)
+
+                if len(files) >= max_results:
+                    break
+
+            return {
+                "success": True,
+                "directory": directory,
+                "pattern": pattern,
+                "files": files,
+                "count": len(files),
+                "truncated": len(files) >= max_results,
+            }
+        except Exception as e:
+            return {"success": False, "error": f"Error finding files: {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("terminal_navigation")
+class TerminalNavigationTool(StatefulTool):
+    """Tool for navigating through the terminal with state memory."""
+
+    def __init__(self):
+        self.name = "terminal_navigation"
+        self.description = (
+            "Navigate through the terminal with persistent directory state"
+        )
+        self.parameters = {
+            "type": "object",
+            "properties": {
+                "action": {
+                    "type": "string",
+                    "enum": ["cd", "pwd", "pushd", "popd", "dirs"],
+                    "description": "Navigation action to perform",
+                },
+                "directory": {
+                    "type": "string",
+                    "description": "Target directory for cd and pushd actions",
+                },
+            },
+            "required": ["action"],
+        }
+        self.reset()
+
+    @classmethod
+    def create_instance(cls):
+        """Create a new instance with fresh state."""
+        return cls()
+
+    def reset(self):
+        """Reset the terminal navigation state."""
+        self.current_dir = os.getcwd()
+        self.dir_stack = []
+
+    def __call__(self, action: str, directory: Optional[str] = None) -> Dict[str, Any]:
+        """Execute the terminal navigation action."""
+        try:
+            # Handle the different navigation actions
+            if action == "cd" and directory:
+                # Expand user path if present
+                target_dir = os.path.expanduser(directory)
+
+                # Handle relative paths
+                if not os.path.isabs(target_dir):
+                    target_dir = os.path.join(self.current_dir, target_dir)
+
+                # Normalize the path
+                target_dir = os.path.normpath(target_dir)
+
+                # Check if directory exists
+                if not os.path.exists(target_dir):
+                    return {
+                        "success": False,
+                        "error": f"Directory does not exist: {target_dir}",
+                        "current_dir": self.current_dir,
+                    }
+
+                if not os.path.isdir(target_dir):
+                    return {
+                        "success": False,
+                        "error": f"Path is not a directory: {target_dir}",
+                        "current_dir": self.current_dir,
+                    }
+
+                # Change directory
+                self.current_dir = target_dir
+                return {
+                    "success": True,
+                    "action": "cd",
+                    "previous_dir": self.current_dir,
+                    "current_dir": self.current_dir,
+                }
+
+            elif action == "pwd":
+                # Print working directory
+                return {
+                    "success": True,
+                    "action": "pwd",
+                    "current_dir": self.current_dir,
+                }
+
+            elif action == "pushd" and directory:
+                # Expand user path if present
+                target_dir = os.path.expanduser(directory)
+
+                # Handle relative paths
+                if not os.path.isabs(target_dir):
+                    target_dir = os.path.join(self.current_dir, target_dir)
+
+                # Normalize the path
+                target_dir = os.path.normpath(target_dir)
+
+                # Check if directory exists
+                if not os.path.exists(target_dir):
+                    return {
+                        "success": False,
+                        "error": f"Directory does not exist: {target_dir}",
+                        "current_dir": self.current_dir,
+                        "dir_stack": self.dir_stack,
+                    }
+
+                if not os.path.isdir(target_dir):
+                    return {
+                        "success": False,
+                        "error": f"Path is not a directory: {target_dir}",
+                        "current_dir": self.current_dir,
+                        "dir_stack": self.dir_stack,
+                    }
+
+                # Push current directory onto stack and change to new directory
+                self.dir_stack.append(self.current_dir)
+                self.current_dir = target_dir
+
+                return {
+                    "success": True,
+                    "action": "pushd",
+                    "previous_dir": self.dir_stack[-1],
+                    "current_dir": self.current_dir,
+                    "dir_stack": self.dir_stack,
+                }
+
+            elif action == "popd":
+                # Check if directory stack is empty
+                if not self.dir_stack:
+                    return {
+                        "success": False,
+                        "error": "Directory stack is empty",
+                        "current_dir": self.current_dir,
+                        "dir_stack": [],
+                    }
+
+                # Pop directory from stack and change to it
+                previous_dir = self.current_dir
+                self.current_dir = self.dir_stack.pop()
+
+                return {
+                    "success": True,
+                    "action": "popd",
+                    "previous_dir": previous_dir,
+                    "current_dir": self.current_dir,
+                    "dir_stack": self.dir_stack,
+                }
+
+            elif action == "dirs":
+                # Display directory stack
+                return {
+                    "success": True,
+                    "action": "dirs",
+                    "current_dir": self.current_dir,
+                    "dir_stack": self.dir_stack,
+                }
+
+            else:
+                return {
+                    "success": False,
+                    "error": f"Invalid action '{action}' or missing required parameters",
+                    "current_dir": self.current_dir,
+                }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"Error navigating terminal: {str(e)}",
+                "current_dir": self.current_dir,
+            }
+
+    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,
+            },
+        }