Jiayi Zhang
d0cf79b6f4
Pre-commit checks / pre-commit-check (push) Failing after 1s
Add E2B sandbox for code execution, and basic Sandbox for computer use.
119 lines
5.0 KiB
Python
119 lines
5.0 KiB
Python
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
from typing import Dict
|
|
from app.tool.base import BaseTool
|
|
import asyncio
|
|
import multiprocessing
|
|
import sys
|
|
from io import StringIO
|
|
from app.config import config
|
|
|
|
class SandboxPythonExecute(BaseTool):
|
|
"""A tool for executing Python code either locally (with timeout) or in a sandboxed environment using e2b_code_interpreter."""
|
|
|
|
name: str = "python_execute"
|
|
description: str = (
|
|
"Executes Python code string. Note: Only print outputs are visible, function return values are not captured. Use print statements to see results. "
|
|
"Set mode='sandbox' to run in a secure sandbox (e2b), otherwise runs locally."
|
|
)
|
|
parameters: dict = {
|
|
"type": "object",
|
|
"properties": {
|
|
"code": {
|
|
"type": "string",
|
|
"description": "The Python code to execute.",
|
|
},
|
|
"mode": {
|
|
"type": "string",
|
|
"enum": ["local", "sandbox"],
|
|
"description": "Execution mode: 'local' (default) or 'sandbox' (e2b sandbox)",
|
|
},
|
|
"timeout": {
|
|
"type": "integer",
|
|
"description": "Execution timeout in seconds (default: 5 for local, 10 for sandbox)",
|
|
},
|
|
},
|
|
"required": ["code"],
|
|
}
|
|
|
|
def _run_code(self, code: str, result_dict: dict, safe_globals: dict) -> None:
|
|
original_stdout = sys.stdout
|
|
try:
|
|
output_buffer = StringIO()
|
|
sys.stdout = output_buffer
|
|
exec(code, safe_globals, safe_globals)
|
|
result_dict["observation"] = output_buffer.getvalue()
|
|
result_dict["success"] = True
|
|
except Exception as e:
|
|
result_dict["observation"] = str(e)
|
|
result_dict["success"] = False
|
|
finally:
|
|
sys.stdout = original_stdout
|
|
|
|
async def execute(
|
|
self,
|
|
code: str,
|
|
timeout: int = None,
|
|
mode: str = "local",
|
|
) -> Dict:
|
|
"""
|
|
Executes the provided Python code in the selected environment.
|
|
Args:
|
|
code (str): The Python code to execute.
|
|
timeout (int): Execution timeout in seconds.
|
|
mode (str): 'local' or 'sandbox'.
|
|
Returns:
|
|
Dict: Contains 'observation' with execution output or error message and 'success' status.
|
|
"""
|
|
if mode == "sandbox":
|
|
# Use e2b_code_interpreter Sandbox
|
|
try:
|
|
from e2b_code_interpreter import Sandbox
|
|
except ImportError:
|
|
return {"observation": "e2b_code_interpreter not installed.", "success": False}
|
|
# Get sandbox config from app.config.config
|
|
sandbox_config = getattr(config, "cloud_sandbox", {})
|
|
# You can pass config to Sandbox() if needed, e.g. Sandbox(api_key=sandbox_config.get("api_key"))
|
|
sbx = Sandbox(**sandbox_config) if sandbox_config else Sandbox()
|
|
try:
|
|
# Default timeout for sandbox is 10s if not set
|
|
effective_timeout = timeout if timeout is not None else 10
|
|
async def run():
|
|
execution = sbx.run_code(code)
|
|
logs = execution.logs if hasattr(execution, 'logs') else str(execution)
|
|
success = execution.error is None if hasattr(execution, 'error') else True
|
|
observation = logs
|
|
if not success:
|
|
observation = f"Error: {getattr(execution, 'error', 'Unknown error')}\n{logs}"
|
|
return {"observation": observation, "success": success}
|
|
result = await asyncio.wait_for(run(), timeout=effective_timeout)
|
|
return result
|
|
except asyncio.TimeoutError:
|
|
return {"observation": f"Execution timeout after {effective_timeout} seconds", "success": False}
|
|
except Exception as e:
|
|
return {"observation": str(e), "success": False}
|
|
finally:
|
|
sbx.kill()
|
|
else:
|
|
# Local execution (default)
|
|
effective_timeout = timeout if timeout is not None else 5
|
|
with multiprocessing.Manager() as manager:
|
|
result = manager.dict({"observation": "", "success": False})
|
|
if isinstance(__builtins__, dict):
|
|
safe_globals = {"__builtins__": __builtins__}
|
|
else:
|
|
safe_globals = {"__builtins__": __builtins__.__dict__.copy()}
|
|
proc = multiprocessing.Process(
|
|
target=self._run_code, args=(code, result, safe_globals)
|
|
)
|
|
proc.start()
|
|
proc.join(effective_timeout)
|
|
if proc.is_alive():
|
|
proc.terminate()
|
|
proc.join(1)
|
|
return {
|
|
"observation": f"Execution timeout after {effective_timeout} seconds",
|
|
"success": False,
|
|
}
|
|
return dict(result)
|