alita-sdk 0.3.379__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/runtime/tools/sandbox.py

@@ -2,9 +2,12 @@ import asyncio
 import logging
 import subprocess
 import os
-from typing import Any, Type, Optional, Dict, List, Literal
+from typing import Any, Type, Optional, Dict, List, Literal, Union
+from copy import deepcopy
+from pathlib import Path
 
 from langchain_core.tools import BaseTool, BaseToolkit
+from langchain_core.messages import ToolCall
 from pydantic import BaseModel, create_model, ConfigDict, Field
 from pydantic.fields import FieldInfo
 
@@ -19,7 +22,7 @@ def get_tools(tools_list: list, alita_client=None, llm=None, memory_store=None):
 
     Args:
         tools_list: List of tool configurations
-        alita_client: Alita client instance (unused for sandbox)
+        alita_client: Alita client instance for sandbox tools
         llm: LLM client instance (unused for sandbox)
         memory_store: Optional memory store instance (unused for sandbox)
 
@@ -34,6 +37,7 @@ def get_tools(tools_list: list, alita_client=None, llm=None, memory_store=None):
                 toolkit_instance = SandboxToolkit.get_toolkit(
                     stateful=tool['settings'].get('stateful', False),
                     allow_net=tool['settings'].get('allow_net', True),
+                    alita_client=alita_client,
                     toolkit_name=tool.get('toolkit_name', '')
                 )
                 all_tools.extend(toolkit_instance.get_tools())
@@ -60,36 +64,10 @@ def _is_deno_available() -> bool:
 
 
 def _setup_pyodide_cache_env() -> None:
-    """Setup Pyodide caching environment variables for performance optimization"""
+    """Setup Pyodide caching environment variables for performance optimization [NO-OP]"""
     try:
-        # Check if cache environment file exists and source it
-        cache_env_file = os.path.expanduser("~/.pyodide_cache_env")
-        if os.path.exists(cache_env_file):
-            with open(cache_env_file, 'r') as f:
-                for line in f:
-                    line = line.strip()
-                    if line.startswith('export ') and '=' in line:
-                        # Parse export VAR=value format
-                        var_assignment = line[7:]  # Remove 'export '
-                        if '=' in var_assignment:
-                            key, value = var_assignment.split('=', 1)
-                            # Remove quotes if present
-                            value = value.strip('"').strip("'")
-                            os.environ[key] = value
-                            logger.debug(f"Set Pyodide cache env: {key}={value}")
-
-        # Set default caching environment variables if not already set
-        cache_defaults = {
-            'PYODIDE_PACKAGES_PATH': os.path.expanduser('~/.cache/pyodide'),
-            'DENO_DIR': os.path.expanduser('~/.cache/deno'),
-            'PYODIDE_CACHE_DIR': os.path.expanduser('~/.cache/pyodide'),
-        }
-
-        for key, default_value in cache_defaults.items():
-            if key not in os.environ:
-                os.environ[key] = default_value
-                logger.debug(f"Set default Pyodide env: {key}={default_value}")
-
+        for key in ["SANDBOX_BASE", "DENO_DIR"]:
+            logger.info("Sandbox env: %s -> %s", key, os.environ.get(key, "n/a"))
     except Exception as e:
         logger.warning(f"Could not setup Pyodide cache environment: {e}")
 
@@ -126,6 +104,7 @@ class PyodideSandboxTool(BaseTool):
     allow_net: bool = True
     session_bytes: Optional[bytes] = None
     session_metadata: Optional[Dict] = None
+    alita_client: Optional[Any] = None
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
@@ -134,6 +113,28 @@ class PyodideSandboxTool(BaseTool):
         _setup_pyodide_cache_env()
         self._initialize_sandbox()
 
+    def _prepare_pyodide_input(self, code: str) -> str:
+        """Prepare input for PyodideSandboxTool by injecting state and alita_client into the code block."""
+        pyodide_predata = ""
+
+        # Add alita_client if available
+        if self.alita_client:
+            try:
+                # Get the directory of the current file and construct the path to sandbox_client.py
+                current_dir = Path(__file__).parent
+                sandbox_client_path = current_dir.parent / 'clients' / 'sandbox_client.py'
+
+                with open(sandbox_client_path, 'r') as f:
+                    sandbox_client_code = f.read()
+                pyodide_predata += f"{sandbox_client_code}\n"
+                pyodide_predata += (f"alita_client = SandboxClient(base_url='{self.alita_client.base_url}',"
+                                    f"project_id={self.alita_client.project_id},"
+                                    f"auth_token='{self.alita_client.auth_token}')\n")
+            except FileNotFoundError:
+                logger.error(f"sandbox_client.py not found. Ensure the file exists.")
+
+        return f"#elitea simplified client\n{pyodide_predata}{code}"
+
     def _initialize_sandbox(self) -> None:
         """Initialize the PyodideSandbox instance with optimized settings"""
         try:
@@ -148,9 +149,19 @@ class PyodideSandboxTool(BaseTool):
 
             from langchain_sandbox import PyodideSandbox
 
+            # Air-gapped settings
+            sandbox_base = os.environ.get("SANDBOX_BASE", os.path.expanduser('~/.cache/pyodide'))
+            sandbox_tmp = os.path.join(sandbox_base, "tmp")
+            deno_cache = os.environ.get("DENO_DIR", os.path.expanduser('~/.cache/deno'))
+
             # Configure sandbox with performance optimizations
             self._sandbox = PyodideSandbox(
                 stateful=self.stateful,
+                #
+                allow_env=["SANDBOX_BASE"],
+                allow_read=[sandbox_base, sandbox_tmp, deno_cache],
+                allow_write=[sandbox_tmp, deno_cache],
+                #
                 allow_net=self.allow_net,
                 # Use auto node_modules_dir for better caching
                 node_modules_dir="auto"
@@ -180,6 +191,9 @@ class PyodideSandboxTool(BaseTool):
             if self._sandbox is None:
                 self._initialize_sandbox()
 
+            # Prepare code with state and client injection
+            prepared_code = self._prepare_pyodide_input(code)
+
             # Check if we're already in an async context
             try:
                 loop = asyncio.get_running_loop()
@@ -187,11 +201,11 @@ class PyodideSandboxTool(BaseTool):
                 # We'll need to use a different approach
                 import concurrent.futures
                 with concurrent.futures.ThreadPoolExecutor() as executor:
-                    future = executor.submit(asyncio.run, self._arun(code))
+                    future = executor.submit(asyncio.run, self._arun(prepared_code))
                     return future.result()
             except RuntimeError:
                 # No running loop, safe to use asyncio.run
-                return asyncio.run(self._arun(code))
+                return asyncio.run(self._arun(prepared_code))
         except (ImportError, RuntimeError) as e:
             # Handle specific dependency errors gracefully
             error_msg = str(e)
@@ -250,7 +264,7 @@ class PyodideSandboxTool(BaseTool):
 
         except Exception as e:
             logger.error(f"Error executing code in sandbox: {e}")
-            return f"Error executing code: {str(e)}"
+            return {"error": f"Error executing code: {str(e)}"}
 
 
 class StatefulPyodideSandboxTool(PyodideSandboxTool):
@@ -278,7 +292,7 @@ class StatefulPyodideSandboxTool(PyodideSandboxTool):
 
 
 # Factory function for creating sandbox tools
-def create_sandbox_tool(stateful: bool = False, allow_net: bool = True) -> BaseTool:
+def create_sandbox_tool(stateful: bool = False, allow_net: bool = True, alita_client: Optional[Any] = None) -> BaseTool:
     """
     Factory function to create sandbox tools with specified configuration.
 
@@ -302,22 +316,25 @@ def create_sandbox_tool(stateful: bool = False, allow_net: bool = True) -> BaseT
         - Cached wheels reduce package download time from ~4.76s to near-instant
     """
     if stateful:
-        return StatefulPyodideSandboxTool(allow_net=allow_net)
+        return StatefulPyodideSandboxTool(allow_net=allow_net, alita_client=alita_client)
     else:
-        return PyodideSandboxTool(stateful=False, allow_net=allow_net)
+        return PyodideSandboxTool(stateful=False, allow_net=allow_net, alita_client=alita_client)
 
 
 class SandboxToolkit(BaseToolkit):
     tools: List[BaseTool] = []
 
     @staticmethod
-    def toolkit_config_schema() -> BaseModel:
-        # Create sample tools to get their schemas
-        sample_tools = [
-            PyodideSandboxTool(),
-            StatefulPyodideSandboxTool()
-        ]
-        selected_tools = {x.name: x.args_schema.schema() for x in sample_tools}
+    def toolkit_config_schema() -> Type[BaseModel]:
+        # Get tool schemas without instantiating the tools (avoids Deno requirement)
+        try:
+            selected_tools = {
+                "pyodide_sandbox": sandbox_tool_input.model_json_schema(),
+                "stateful_pyodide_sandbox": sandbox_tool_input.model_json_schema(),
+            }
+        except Exception as e:
+            logger.warning(f"Could not generate sandbox tool schemas: {e}")
+            selected_tools = {}
 
         return create_model(
             'sandbox',
@@ -338,24 +355,24 @@ class SandboxToolkit(BaseToolkit):
         )
 
     @classmethod
-    def get_toolkit(cls, stateful: bool = False, allow_net: bool = True, **kwargs):
+    def get_toolkit(cls, stateful: bool = False, allow_net: bool = True, alita_client=None, **kwargs):
         """
         Get toolkit with sandbox tools.
 
         Args:
             stateful: Whether to maintain state between executions
             allow_net: Whether to allow network access
+            alita_client: Alita client instance for sandbox tools
             **kwargs: Additional arguments
         """
         tools = []
 
         if stateful:
-            tools.append(StatefulPyodideSandboxTool(allow_net=allow_net))
+            tools.append(StatefulPyodideSandboxTool(allow_net=allow_net, alita_client=alita_client))
         else:
-            tools.append(PyodideSandboxTool(stateful=False, allow_net=allow_net))
+            tools.append(PyodideSandboxTool(stateful=False, allow_net=allow_net, alita_client=alita_client))
 
         return cls(tools=tools)
 
     def get_tools(self):
         return self.tools
-