ai-computer-client 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

ai_computer/models.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+from typing import Optional, Dict
+
+@dataclass
+class SandboxResponse:
+    """Response from sandbox operations.
+    
+    Attributes:
+        success: Whether the operation was successful
+        data: Optional response data
+        error: Optional error message if operation failed
+    """
+    success: bool
+    data: Optional[Dict] = None
+    error: Optional[str] = None
+
+@dataclass
+class StreamEvent:
+    """Event from streaming code execution.
+    
+    Attributes:
+        type: Type of event ('stdout', 'stderr', 'info', 'error', 'completed', 'keepalive')
+        data: Event data
+    """
+    type: str
+    data: str
+
+@dataclass
+class FileOperationResponse:
+    """Response from file operations.
+    
+    Attributes:
+        success: Whether the operation was successful
+        filename: Name of the file
+        size: Size of the file in bytes
+        path: Path where the file was saved
+        message: Optional status message
+        error: Optional error message if operation failed
+    """
+    success: bool
+    filename: Optional[str] = None
+    size: Optional[int] = None
+    path: Optional[str] = None
+    message: Optional[str] = None
+    error: Optional[str] = None

ai_computer/submodules/__init__.py

@@ -0,0 +1,5 @@
+from .filesystem import FileSystemModule
+from .shell import ShellModule
+from .code import CodeModule
+
+__all__ = ["FileSystemModule", "ShellModule", "CodeModule"]

ai_computer/submodules/base.py

@@ -0,0 +1,81 @@
+from typing import Optional, Dict, Any
+import aiohttp
+from ..models import SandboxResponse
+
+class BaseSubmodule:
+    """Base class for all submodules.
+    
+    This class provides common functionality for all submodules, including
+    access to the parent client's authentication token and sandbox ID.
+    
+    Attributes:
+        _client: Reference to the parent SandboxClient
+    """
+    
+    def __init__(self, client):
+        """Initialize the submodule.
+        
+        Args:
+            client: The parent SandboxClient instance
+        """
+        self._client = client
+    
+    @property
+    def base_url(self) -> str:
+        """Get the base URL from the parent client."""
+        return self._client.base_url
+    
+    @property
+    def token(self) -> Optional[str]:
+        """Get the authentication token from the parent client."""
+        return self._client.token
+    
+    @property
+    def sandbox_id(self) -> Optional[str]:
+        """Get the sandbox ID from the parent client."""
+        return self._client.sandbox_id
+    
+    async def _ensure_ready(self) -> SandboxResponse:
+        """Ensure the sandbox is ready for operations.
+        
+        Returns:
+            SandboxResponse indicating if the sandbox is ready
+        """
+        if not self.token or not self.sandbox_id:
+            return SandboxResponse(
+                success=False, 
+                error="Client not properly initialized. Call setup() first"
+            )
+        
+        # Ensure sandbox is ready
+        return await self._client.wait_for_ready()
+    
+    def _get_headers(self, content_type: str = "application/json") -> Dict[str, str]:
+        """Get the headers for API requests.
+        
+        Args:
+            content_type: The content type for the request
+            
+        Returns:
+            Dictionary of headers
+        """
+        return {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": content_type
+        }
+    
+    async def _handle_response(self, response: aiohttp.ClientResponse) -> SandboxResponse:
+        """Handle the API response.
+        
+        Args:
+            response: The aiohttp response object
+            
+        Returns:
+            SandboxResponse with the parsed response data
+        """
+        if response.status != 200:
+            error_text = await response.text()
+            return SandboxResponse(success=False, error=error_text)
+        
+        result = await response.json()
+        return SandboxResponse(success=True, data=result)

ai_computer/submodules/code.py

@@ -0,0 +1,295 @@
+import aiohttp
+import asyncio
+from typing import Optional, Dict, AsyncGenerator, Any
+
+from .base import BaseSubmodule
+from ..models import SandboxResponse, StreamEvent
+
+class CodeModule(BaseSubmodule):
+    """Code execution operations for the sandbox environment.
+    
+    This module provides methods for executing Python code in the sandbox.
+    """
+    
+    async def execute(
+        self,
+        code: str,
+        timeout: int = 30,
+        environment: Optional[Dict[str, str]] = None
+    ) -> SandboxResponse:
+        """Execute Python code in the sandbox.
+        
+        Args:
+            code: The Python code to execute
+            timeout: Maximum execution time in seconds
+            environment: Optional environment variables for the execution
+            
+        Returns:
+            SandboxResponse containing execution results
+        """
+        ready = await self._ensure_ready()
+        if not ready.success:
+            return ready
+            
+        headers = self._get_headers()
+        
+        data = {
+            "code": code,
+            "timeout": timeout
+        }
+        
+        if environment:
+            data["environment"] = environment
+        
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute",
+                    headers=headers,
+                    json=data
+                ) as response:
+                    return await self._handle_response(response)
+                    
+        except Exception as e:
+            return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
+    
+    async def execute_stream(
+        self,
+        code: str,
+        timeout: int = 30,
+        environment: Optional[Dict[str, str]] = None
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Execute Python code in the sandbox with streaming output.
+        
+        Args:
+            code: The Python code to execute
+            timeout: Maximum execution time in seconds
+            environment: Optional environment variables for the execution
+            
+        Yields:
+            StreamEvent objects containing execution output
+        """
+        ready = await self._ensure_ready()
+        if not ready.success:
+            yield StreamEvent(type="error", data=ready.error or "Sandbox not ready")
+            return
+            
+        headers = self._get_headers()
+        
+        data = {
+            "code": code,
+            "timeout": timeout
+        }
+        
+        if environment:
+            data["environment"] = environment
+        
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute/stream",
+                    headers=headers,
+                    json=data
+                ) as response:
+                    if response.status != 200:
+                        error_text = await response.text()
+                        yield StreamEvent(type="error", data=error_text)
+                        return
+                    
+                    # Process the streaming response
+                    async for line in response.content:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        
+                        try:
+                            event_data = line.decode('utf-8')
+                            # Check if it's a JSON object
+                            if event_data.startswith('{') and event_data.endswith('}'):
+                                import json
+                                event_json = json.loads(event_data)
+                                event_type = event_json.get('type', 'info')
+                                event_data = event_json.get('data', '')
+                                yield StreamEvent(type=event_type, data=event_data)
+                                
+                                # If execution is complete, break the loop
+                                if event_type == 'completed':
+                                    break
+                            else:
+                                # Treat as stdout if not JSON
+                                yield StreamEvent(type="stdout", data=event_data)
+                        except Exception as e:
+                            yield StreamEvent(type="error", data=f"Failed to parse event: {str(e)}")
+                            break
+                                
+        except Exception as e:
+            yield StreamEvent(type="error", data=f"Connection error: {str(e)}")
+    
+    async def execute_file(
+        self,
+        file_path: str,
+        timeout: int = 30,
+        environment: Optional[Dict[str, str]] = None
+    ) -> SandboxResponse:
+        """Execute a Python file in the sandbox.
+        
+        Args:
+            file_path: Path to the Python file in the sandbox
+            timeout: Maximum execution time in seconds
+            environment: Optional environment variables for the execution
+            
+        Returns:
+            SandboxResponse containing execution results
+        """
+        ready = await self._ensure_ready()
+        if not ready.success:
+            return ready
+            
+        # Execute code to check if the file exists and is a Python file
+        check_code = f"""
+import os
+
+file_path = "{file_path}"
+if not os.path.exists(file_path):
+    result = {{"success": False, "error": f"File not found: {{file_path}}"}}
+elif not os.path.isfile(file_path):
+    result = {{"success": False, "error": f"Not a file: {{file_path}}"}}
+elif not file_path.endswith('.py'):
+    result = {{"success": False, "error": f"Not a Python file: {{file_path}}"}}
+else:
+    result = {{"success": True}}
+
+result
+"""
+        check_response = await self.execute(check_code)
+        if not check_response.success:
+            return check_response
+            
+        result = check_response.data.get('result', {})
+        if not result.get('success', False):
+            return SandboxResponse(
+                success=False,
+                error=result.get('error', 'Unknown error checking file')
+            )
+            
+        # Execute the Python file
+        headers = self._get_headers()
+        
+        data = {
+            "file_path": file_path,
+            "timeout": timeout
+        }
+        
+        if environment:
+            data["environment"] = environment
+        
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute",
+                    headers=headers,
+                    json=data
+                ) as response:
+                    return await self._handle_response(response)
+                    
+        except Exception as e:
+            return SandboxResponse(success=False, error=f"Connection error: {str(e)}")
+    
+    async def execute_file_stream(
+        self,
+        file_path: str,
+        timeout: int = 30,
+        environment: Optional[Dict[str, str]] = None
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Execute a Python file in the sandbox with streaming output.
+        
+        Args:
+            file_path: Path to the Python file in the sandbox
+            timeout: Maximum execution time in seconds
+            environment: Optional environment variables for the execution
+            
+        Yields:
+            StreamEvent objects containing execution output
+        """
+        ready = await self._ensure_ready()
+        if not ready.success:
+            yield StreamEvent(type="error", data=ready.error or "Sandbox not ready")
+            return
+            
+        # Execute code to check if the file exists and is a Python file
+        check_code = f"""
+import os
+
+file_path = "{file_path}"
+if not os.path.exists(file_path):
+    result = {{"success": False, "error": f"File not found: {{file_path}}"}}
+elif not os.path.isfile(file_path):
+    result = {{"success": False, "error": f"Not a file: {{file_path}}"}}
+elif not file_path.endswith('.py'):
+    result = {{"success": False, "error": f"Not a Python file: {{file_path}}"}}
+else:
+    result = {{"success": True}}
+
+result
+"""
+        check_response = await self.execute(check_code)
+        if not check_response.success:
+            yield StreamEvent(type="error", data=check_response.error or "Failed to check file")
+            return
+            
+        result = check_response.data.get('result', {})
+        if not result.get('success', False):
+            yield StreamEvent(type="error", data=result.get('error', 'Unknown error checking file'))
+            return
+            
+        # Execute the Python file with streaming output
+        headers = self._get_headers()
+        
+        data = {
+            "file_path": file_path,
+            "timeout": timeout
+        }
+        
+        if environment:
+            data["environment"] = environment
+        
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    f"{self.base_url}/api/v1/sandbox/{self.sandbox_id}/execute/file/stream",
+                    headers=headers,
+                    json=data
+                ) as response:
+                    if response.status != 200:
+                        error_text = await response.text()
+                        yield StreamEvent(type="error", data=error_text)
+                        return
+                    
+                    # Process the streaming response
+                    async for line in response.content:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        
+                        try:
+                            event_data = line.decode('utf-8')
+                            # Check if it's a JSON object
+                            if event_data.startswith('{') and event_data.endswith('}'):
+                                import json
+                                event_json = json.loads(event_data)
+                                event_type = event_json.get('type', 'info')
+                                event_data = event_json.get('data', '')
+                                yield StreamEvent(type=event_type, data=event_data)
+                                
+                                # If execution is complete, break the loop
+                                if event_type == 'completed':
+                                    break
+                            else:
+                                # Treat as stdout if not JSON
+                                yield StreamEvent(type="stdout", data=event_data)
+                        except Exception as e:
+                            yield StreamEvent(type="error", data=f"Failed to parse event: {str(e)}")
+                            break
+                                
+        except Exception as e:
+            yield StreamEvent(type="error", data=f"Connection error: {str(e)}")