#!/usr/bin/env python # coding=utf-8 # Copyright 2024 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import base64 import inspect import json import os import pickle import re import secrets import subprocess import tempfile import time import uuid from contextlib import closing from io import BytesIO from textwrap import dedent from typing import Any, Optional import PIL.Image import requests from requests.exceptions import RequestException from .default_tools import FinalAnswerTool from .local_python_executor import CodeOutput, PythonExecutor from .monitoring import LogLevel from .serialization import SafeSerializer, SerializationError from .tools import Tool, get_tools_definition_code from .utils import AgentError __all__ = ["BlaxelExecutor", "E2BExecutor", "ModalExecutor", "DockerExecutor", "WasmExecutor"] try: from dotenv import load_dotenv load_dotenv() except ModuleNotFoundError: pass class RemotePythonExecutor(PythonExecutor): """ Executor of Python code in a remote environment. Args: additional_imports (`list[str]`): Additional Python packages to install. logger (`Logger`): Logger to use for output and errors. allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` if you fully trust the execution environment and need backward compatibility with custom types. """ FINAL_ANSWER_EXCEPTION = "FinalAnswerException" def __init__( self, additional_imports: list[str], logger, allow_pickle: bool = False, ): self.additional_imports = additional_imports self.logger = logger self.allow_pickle = allow_pickle self.logger.log("Initializing executor, hold on...") self.installed_packages = [] def run_code_raise_errors(self, code: str) -> CodeOutput: """ Execute Python code in the remote environment and return the result. Args: code (`str`): Python code to execute. Returns: `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. """ raise NotImplementedError def send_tools(self, tools: dict[str, Tool]): if "final_answer" in tools: self._patch_final_answer_with_exception(tools["final_answer"]) # Install tool packages packages_to_install = { pkg for tool in tools.values() for pkg in tool.to_dict()["requirements"] if pkg not in self.installed_packages + ["smolagents"] } if "PIL" in packages_to_install: packages_to_install.discard("PIL") packages_to_install.add("pillow") if packages_to_install: self.installed_packages += self.install_packages(list(packages_to_install)) # Get tool definitions code = get_tools_definition_code(tools) if code: code_output = self.run_code_raise_errors(code) self.logger.log(code_output.logs) def send_variables(self, variables: dict[str, Any]): """Send variables to the kernel namespace using SafeSerializer. Uses prefix-based format ("safe:..." or "pickle:..."). When allow_pickle=False, only safe JSON serialization is allowed. When allow_pickle=True, pickle fallback is enabled for complex types. """ if not variables: return serialized = SafeSerializer.dumps(variables, allow_pickle=self.allow_pickle) code = f""" {SafeSerializer.get_deserializer_code(self.allow_pickle)} vars_dict = _deserialize({repr(serialized)}) locals().update(vars_dict) """ self.run_code_raise_errors(code) def __call__(self, code_action: str) -> CodeOutput: """Run the code and determine if it is the final answer.""" return self.run_code_raise_errors(code_action) def install_packages(self, additional_imports: list[str]): if additional_imports: code_output = self.run_code_raise_errors(f"!pip install {' '.join(additional_imports)}") self.logger.log(code_output.logs) return additional_imports def _patch_final_answer_with_exception(self, final_answer_tool: FinalAnswerTool): """Patch the FinalAnswerTool to raise an exception. This is necessary because the remote executors rely on the FinalAnswerTool to detect the final answer. It modifies the `forward` method of the FinalAnswerTool to raise a `FinalAnswerException` with the final answer as a serialized value. This allows the executor to catch this exception and return the final answer. Uses prefix-based format ("safe:" or "pickle:") for serialization. Args: final_answer_tool (`FinalAnswerTool`): FinalAnswerTool instance to patch. """ # Create a new class that inherits from the original FinalAnswerTool class _FinalAnswerTool(final_answer_tool.__class__): pass # Add a new forward method that raises the FinalAnswerException # NOTE: Serialization logic is inlined here because this method's source code # is extracted and sent to remote environments where external references don't exist # Capture settings via closure allow_pickle_setting = self.allow_pickle def forward(self, *args, **kwargs) -> Any: import base64 import json from io import BytesIO # Baked in from closure at patch time ALLOW_PICKLE = allow_pickle_setting class SerializationError(Exception): pass def _to_json_safe(obj): if isinstance(obj, (str, int, float, bool, type(None))): return obj elif isinstance(obj, dict): # Check if all keys are strings (JSON-compatible) if all(isinstance(k, str) for k in obj.keys()): return {k: _to_json_safe(v) for k, v in obj.items()} else: return { "__type__": "dict_with_complex_keys", "data": [[_to_json_safe(k), _to_json_safe(v)] for k, v in obj.items()], } elif isinstance(obj, list): return [_to_json_safe(item) for item in obj] elif isinstance(obj, tuple): return {"__type__": "tuple", "data": [_to_json_safe(item) for item in obj]} elif isinstance(obj, set): return {"__type__": "set", "data": [_to_json_safe(item) for item in obj]} elif isinstance(obj, bytes): return {"__type__": "bytes", "data": base64.b64encode(obj).decode()} elif isinstance(obj, complex): return {"__type__": "complex", "real": obj.real, "imag": obj.imag} elif isinstance(obj, frozenset): return {"__type__": "frozenset", "data": [_to_json_safe(item) for item in obj]} # Try PIL Image try: import PIL.Image if isinstance(obj, PIL.Image.Image): buffer = BytesIO() obj.save(buffer, format="PNG") return {"__type__": "PIL.Image", "data": base64.b64encode(buffer.getvalue()).decode()} except ImportError: pass # Lazy imports for less common types from datetime import date, datetime, time, timedelta from decimal import Decimal from pathlib import Path if isinstance(obj, datetime): return {"__type__": "datetime", "data": obj.isoformat()} elif isinstance(obj, date): return {"__type__": "date", "data": obj.isoformat()} elif isinstance(obj, time): return {"__type__": "time", "data": obj.isoformat()} elif isinstance(obj, timedelta): return {"__type__": "timedelta", "total_seconds": obj.total_seconds()} elif isinstance(obj, Decimal): return {"__type__": "Decimal", "data": str(obj)} elif isinstance(obj, Path): return {"__type__": "Path", "data": str(obj)} # Try numpy if available try: import numpy as np if isinstance(obj, np.ndarray): return {"__type__": "ndarray", "data": obj.tolist(), "dtype": str(obj.dtype)} elif isinstance(obj, (np.integer, np.floating)): return obj.item() except ImportError: pass # Try dataclass import dataclasses if dataclasses.is_dataclass(obj) and not isinstance(obj, type): return { "__type__": "dataclass", "class_name": type(obj).__name__, "module": type(obj).__module__, "data": {f.name: _to_json_safe(getattr(obj, f.name)) for f in dataclasses.fields(obj)}, } # Cannot safely serialize - raise error for safe mode raise SerializationError(f"Cannot safely serialize object of type {type(obj).__name__}") def _serialize_with_fallback(obj): """Serialize with safe method, fallback to pickle if allowed.""" import pickle if not ALLOW_PICKLE: # Safe ONLY mode - NO pickle fallback, raise error if can't serialize json_safe = _to_json_safe(obj) # Will raise SerializationError if fails return "safe:" + json.dumps(json_safe) else: # Try safe first, fallback to pickle if allowed try: json_safe = _to_json_safe(obj) return "safe:" + json.dumps(json_safe) except SerializationError: # Fallback to pickle try: return "pickle:" + base64.b64encode(pickle.dumps(obj)).decode() except (pickle.PicklingError, TypeError, AttributeError): # Last resort: string representation return "safe:" + json.dumps(str(obj)) class FinalAnswerException(BaseException): def __init__(self, value): self.value = value raise FinalAnswerException(_serialize_with_fallback(self._forward(*args, **kwargs))) # - Set the new forward method function to the _FinalAnswerTool class _FinalAnswerTool.forward = forward # Set __source__ with the actual values baked in (closures don't survive source extraction) source = inspect.getsource(forward) source = source.replace("ALLOW_PICKLE = allow_pickle_setting", f"ALLOW_PICKLE = {allow_pickle_setting}") forward.__source__ = source # Rename the original forward method to _forward # - Get the original forward method function from the final_answer_tool instance original_forward_function = final_answer_tool.forward.__func__ # - Set the new _forward method function to the _FinalAnswerTool class _FinalAnswerTool._forward = original_forward_function # - Update the source code of the new forward method to match the original but with the new name _FinalAnswerTool._forward.__source__ = inspect.getsource(original_forward_function).replace( "def forward(", "def _forward(" ) # Set the new class as the class of the final_answer_tool instance final_answer_tool.__class__ = _FinalAnswerTool @staticmethod def _deserialize_final_answer(encoded_value: str, allow_pickle: bool = True) -> Any: """Deserialize final answer with format detection. Accepts explicit prefix-based formats only: - "safe:" for JSON-safe payloads - "pickle:" for pickle payloads (only when allow_pickle=True) Args: encoded_value: Serialized string from FinalAnswerException. allow_pickle: Whether to allow pickle deserialization. Returns: Deserialized Python object. Raises: SerializationError: If pickle data is rejected. """ if encoded_value.startswith("safe:"): json_data = json.loads(encoded_value[5:]) return SafeSerializer.from_json_safe(json_data) elif encoded_value.startswith("pickle:"): if not allow_pickle: raise SerializationError("Pickle data rejected: allow_pickle=False") return pickle.loads(base64.b64decode(encoded_value[7:])) else: raise SerializationError("Unknown final answer format: expected 'safe:' or 'pickle:' prefix") class E2BExecutor(RemotePythonExecutor): """ Remote Python code executor in an E2B sandbox. Args: additional_imports (`list[str]`): Additional Python packages to install. logger (`Logger`): Logger to use for output and errors. allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` if you fully trust the execution environment and need backward compatibility with custom types. **kwargs: Additional keyword arguments to pass to the E2B Sandbox instantiation. """ def __init__( self, additional_imports: list[str], logger, allow_pickle: bool = False, **kwargs, ): super().__init__(additional_imports, logger, allow_pickle) try: from e2b_code_interpreter import Sandbox except ModuleNotFoundError: raise ModuleNotFoundError( """Please install 'e2b' extra to use E2BExecutor: `pip install 'smolagents[e2b]'`""" ) # Support both e2b v1 and v2 constructors # v2 exposes Sandbox.create(...), while v1 uses Sandbox(...) if hasattr(Sandbox, "create"): self.sandbox = Sandbox.create(**kwargs) else: self.sandbox = Sandbox(**kwargs) self.installed_packages = self.install_packages(additional_imports) self.logger.log("E2B is running", level=LogLevel.INFO) def run_code_raise_errors(self, code: str) -> CodeOutput: """ Execute Python code in the E2B sandbox and return the result. Args: code (`str`): Python code to execute. Returns: `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. """ execution = self.sandbox.run_code(code) execution_logs = "\n".join([str(log) for log in execution.logs.stdout]) # Handle errors if execution.error: # Check if the error is a FinalAnswerException if execution.error.name == RemotePythonExecutor.FINAL_ANSWER_EXCEPTION: final_answer = self._deserialize_final_answer(execution.error.value, self.allow_pickle) return CodeOutput(output=final_answer, logs=execution_logs, is_final_answer=True) # Construct error message error_message = ( f"{execution_logs}\n" f"Executing code yielded an error:\n" f"{execution.error.name}\n" f"{execution.error.value}\n" f"{execution.error.traceback}" ) raise AgentError(error_message, self.logger) # Handle results if not execution.results: return CodeOutput(output=None, logs=execution_logs, is_final_answer=False) for result in execution.results: if not result.is_main_result: continue # Handle image outputs for attribute_name in ["jpeg", "png"]: img_data = getattr(result, attribute_name, None) if img_data is not None: decoded_bytes = base64.b64decode(img_data.encode("utf-8")) return CodeOutput( output=PIL.Image.open(BytesIO(decoded_bytes)), logs=execution_logs, is_final_answer=False ) # Handle other data formats for attribute_name in [ "chart", "data", "html", "javascript", "json", "latex", "markdown", "pdf", "svg", "text", ]: data = getattr(result, attribute_name, None) if data is not None: return CodeOutput(output=data, logs=execution_logs, is_final_answer=False) # If no main result found, return None return CodeOutput(output=None, logs=execution_logs, is_final_answer=False) def cleanup(self): """Clean up the E2B sandbox and resources.""" try: if hasattr(self, "sandbox"): self.logger.log("Shutting down sandbox...", level=LogLevel.INFO) self.sandbox.kill() self.logger.log("Sandbox cleanup completed", level=LogLevel.INFO) del self.sandbox except Exception as e: self.logger.log_error(f"Error during cleanup: {e}") def _websocket_send_execute_request(code: str, ws) -> str: """Send code execution request to kernel.""" import uuid # Generate a unique message ID msg_id = str(uuid.uuid4()) # Create execute request execute_request = { "header": { "msg_id": msg_id, "username": "anonymous", "session": str(uuid.uuid4()), "msg_type": "execute_request", "version": "5.0", }, "parent_header": {}, "metadata": {}, "content": { "code": code, "silent": False, "store_history": True, "user_expressions": {}, "allow_stdin": False, }, } ws.send(json.dumps(execute_request)) return msg_id def _websocket_run_code_raise_errors( code: str, ws, logger, allow_pickle: bool = True, safe_serialization: bool = False ) -> CodeOutput: """Run code over a websocket.""" try: # Send execute request msg_id = _websocket_send_execute_request(code, ws) # Collect output and results outputs = [] result = None is_final_answer = False while True: msg = json.loads(ws.recv()) parent_msg_id = msg.get("parent_header", {}).get("msg_id") # Skip unrelated messages if parent_msg_id != msg_id: continue msg_type = msg.get("msg_type", "") msg_content = msg.get("content", {}) if msg_type == "stream": outputs.append(msg_content["text"]) elif msg_type == "execute_result": result = msg_content["data"].get("text/plain", None) elif msg_type == "error": if msg_content.get("ename", "") == RemotePythonExecutor.FINAL_ANSWER_EXCEPTION: result = RemotePythonExecutor._deserialize_final_answer( msg_content.get("evalue", ""), allow_pickle ) is_final_answer = True else: raise AgentError("\n".join(msg_content.get("traceback", [])), logger) elif msg_type == "status" and msg_content["execution_state"] == "idle": break return CodeOutput(output=result, logs="".join(outputs), is_final_answer=is_final_answer) except Exception as e: logger.log_error(f"Code execution failed: {e}") raise def _create_kernel_http(crate_kernel_endpoint: str, logger, headers: Optional[dict] = None) -> str: """Create kernel using http.""" r = requests.post(crate_kernel_endpoint, headers=headers) if r.status_code != 201: error_details = { "status_code": r.status_code, "headers": dict(r.headers), "url": r.url, "body": r.text, "request_method": r.request.method, "request_headers": dict(r.request.headers), "request_body": r.request.body, } logger.log_error(f"Failed to create kernel. Details: {json.dumps(error_details, indent=2)}") raise RuntimeError(f"Failed to create kernel: Status {r.status_code}\nResponse: {r.text}") from None return r.json()["id"] class DockerExecutor(RemotePythonExecutor): """ Remote Python code executor using Jupyter Kernel Gateway in a Docker container. Args: additional_imports (`list[str]`): Additional Python packages to install. logger (`Logger`): Logger to use for output and errors. allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` if you fully trust the execution environment and need backward compatibility with custom types. host (`str`, default `"127.0.0.1"`): Host to bind to. port (`int`, default `8888`): Port to bind to. image_name (`str`, default `"jupyter-kernel"`): Name of the Docker image to use. If the image doesn't exist, it will be built. build_new_image (`bool`, default `True`): Whether to rebuild a new image even if it already exists. container_run_kwargs (`dict`, *optional*): Additional keyword arguments to pass to the Docker container run command. dockerfile_content (`str`, *optional*): Custom Dockerfile content. If `None`, uses default. """ def __init__( self, additional_imports: list[str], logger, allow_pickle: bool = False, host: str = "127.0.0.1", port: int = 8888, image_name: str = "jupyter-kernel", build_new_image: bool = True, container_run_kwargs: dict[str, Any] | None = None, dockerfile_content: str | None = None, ): super().__init__(additional_imports, logger, allow_pickle) try: import docker except ModuleNotFoundError: raise ModuleNotFoundError( "Please install 'docker' extra to use DockerExecutor: `pip install 'smolagents[docker]'`" ) self.host = host self.port = port self.image_name = image_name self.dockerfile_content = dockerfile_content or dedent( """\ FROM python:3.12-bullseye RUN pip install jupyter_kernel_gateway jupyter_client ipykernel EXPOSE 8888 CMD ["jupyter", "kernelgateway", "--KernelGatewayApp.ip=0.0.0.0", "--KernelGatewayApp.port=8888"] """ ) # Initialize Docker try: self.client = docker.from_env() except docker.errors.DockerException as e: raise RuntimeError("Could not connect to Docker daemon: make sure Docker is running.") from e # Build and start container try: # Check if image exists, unless forced to rebuild if not build_new_image: try: self.client.images.get(self.image_name) self.logger.log(f"Using existing Docker image: {self.image_name}", level=LogLevel.INFO) except docker.errors.ImageNotFound: self.logger.log(f"Image {self.image_name} not found, building...", level=LogLevel.INFO) build_new_image = True if build_new_image: self.logger.log(f"Building Docker image {self.image_name}...", level=LogLevel.INFO) dockerfile_obj = BytesIO(self.dockerfile_content.encode("utf-8")) _, build_logs = self.client.images.build(fileobj=dockerfile_obj, tag=self.image_name) for log_chunk in build_logs: # Only log non-empty messages if log_message := log_chunk.get("stream", "").rstrip(): self.logger.log(log_message, level=LogLevel.DEBUG) self.logger.log(f"Starting container on {host}:{port}...", level=LogLevel.INFO) # Create base container parameters container_kwargs = {} if container_run_kwargs: container_kwargs.update(container_run_kwargs) # Ensure required port mapping and background running if not isinstance(container_kwargs.get("ports"), dict): container_kwargs["ports"] = {} container_kwargs["ports"]["8888/tcp"] = (host, port) container_kwargs["detach"] = True # Generate auth token and pass it to the kernel gateway via the standard KG_AUTH_TOKEN env var token = secrets.token_urlsafe(16) env = container_kwargs.get("environment") or {} if isinstance(env, list): env = dict(kv.split("=", 1) for kv in env if "=" in kv) env["KG_AUTH_TOKEN"] = token container_kwargs["environment"] = env self.container = self.client.containers.run(self.image_name, **container_kwargs) retries = 0 while self.container.status != "running" and retries < 5: self.logger.log(f"Container status: {self.container.status}, waiting...", level=LogLevel.INFO) time.sleep(1) self.container.reload() retries += 1 self.base_url = f"http://{host}:{port}" # Wait for Jupyter to start self._wait_for_server(token) # Create new kernel via HTTP self.kernel_id = _create_kernel_http(f"{self.base_url}/api/kernels?token={token}", self.logger) self.ws_url = f"ws://{host}:{port}/api/kernels/{self.kernel_id}/channels?token={token}" self.installed_packages = self.install_packages(additional_imports) self.logger.log( f"Container {self.container.short_id} is running with kernel {self.kernel_id}", level=LogLevel.INFO ) except Exception as e: self.cleanup() raise RuntimeError(f"Failed to initialize Jupyter kernel: {e}") from e def run_code_raise_errors(self, code: str) -> CodeOutput: """ Execute Python code in the Docker container and return the result. Args: code (`str`): Python code to execute. Returns: `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. """ from websocket import create_connection with closing(create_connection(self.ws_url)) as ws: return _websocket_run_code_raise_errors(code, ws, self.logger, self.allow_pickle) def cleanup(self): """Clean up the Docker container and resources.""" try: if hasattr(self, "container"): self.logger.log(f"Stopping and removing container {self.container.short_id}...", level=LogLevel.INFO) self.container.stop() self.container.remove() self.logger.log("Container cleanup completed", level=LogLevel.INFO) del self.container except Exception as e: self.logger.log_error(f"Error during cleanup: {e}") def delete(self): """Ensure cleanup on deletion.""" self.cleanup() def _wait_for_server(self, token: str): retries = 0 jupyter_ready = False while not jupyter_ready and retries < 10: try: if requests.get(f"{self.base_url}/api/kernelspecs?token={token}", timeout=2).status_code == 200: jupyter_ready = True else: self.logger.log("Jupyter not ready, waiting...", level=LogLevel.INFO) except requests.RequestException: self.logger.log("Jupyter not ready, waiting...", level=LogLevel.INFO) if not jupyter_ready: time.sleep(1) retries += 1 class ModalExecutor(RemotePythonExecutor): """ Remote Python code executor in a Modal sandbox. Args: additional_imports (`list[str]`): Additional Python packages to install. logger (`Logger`): Logger to use for output and errors. allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` if you fully trust the execution environment and need backward compatibility with custom types. app_name (`str`, default `"smolagent-executor"`): App name. port (`int`, default `8888`): Port for jupyter to bind to. create_kwargs (`dict`, *optional*): Additional keyword arguments to pass to the Modal Sandbox create command. See `modal.Sandbox.create` [docs](https://modal.com/docs/reference/modal.Sandbox#create) for all the keyword arguments. """ _ANSI_ESCAPE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") def __init__( self, additional_imports: list[str], logger, allow_pickle: bool = False, app_name: str = "smolagent-executor", port: int = 8888, create_kwargs: Optional[dict] = None, ): super().__init__(additional_imports, logger, allow_pickle) self.port = port try: import modal except ModuleNotFoundError: raise ModuleNotFoundError( """Please install 'modal' extra to use ModalExecutor: `pip install 'smolagents[modal]'`""" ) if create_kwargs is None: create_kwargs = {} create_kwargs = { "image": modal.Image.debian_slim().uv_pip_install("jupyter_kernel_gateway", "ipykernel"), "timeout": 60 * 5, **create_kwargs, } if "app" not in create_kwargs: create_kwargs["app"] = modal.App.lookup(app_name, create_if_missing=True) if "encrypted_ports" not in create_kwargs: create_kwargs["encrypted_ports"] = [port] else: create_kwargs["encrypted_ports"] = create_kwargs["encrypted_ports"] + [port] token = secrets.token_urlsafe(16) default_secrets = [modal.Secret.from_dict({"KG_AUTH_TOKEN": token})] if "secrets" not in create_kwargs: create_kwargs["secrets"] = default_secrets else: create_kwargs["secrets"] = create_kwargs["secrets"] + default_secrets entrypoint = [ "jupyter", "kernelgateway", "--KernelGatewayApp.ip=0.0.0.0", f"--KernelGatewayApp.port={port}", ] self.logger.log("Starting Modal sandbox", level=LogLevel.INFO) self.sandbox = modal.Sandbox.create( *entrypoint, **create_kwargs, ) tunnel = self.sandbox.tunnels()[port] self.logger.log(f"Waiting for Modal sandbox on {tunnel.host}:{port}", level=LogLevel.INFO) self._wait_for_server(tunnel.host, token) self.logger.log("Starting Jupyter kernel", level=LogLevel.INFO) kernel_id = _create_kernel_http(f"https://{tunnel.host}/api/kernels?token={token}", logger) self.ws_url = f"wss://{tunnel.host}/api/kernels/{kernel_id}/channels?token={token}" self.installed_packages = self.install_packages(additional_imports) def run_code_raise_errors(self, code: str) -> CodeOutput: """ Execute Python code in the Modal sandbox and return the result. Args: code (`str`): Python code to execute. Returns: `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. """ from websocket import create_connection with closing(create_connection(self.ws_url)) as ws: return _websocket_run_code_raise_errors(code, ws, self.logger, self.allow_pickle) def cleanup(self): """Clean up the Modal sandbox by terminating it.""" if hasattr(self, "sandbox"): self.sandbox.terminate() def delete(self): """Ensure cleanup on deletion.""" self.cleanup() def _wait_for_server(self, host: str, token: str): """Wait for server to start up.""" n_retries = 0 while True: try: resp = requests.get(f"https://{host}/api/kernelspecs?token={token}") if resp.status_code == 200: break except RequestException: n_retries += 1 if n_retries % 10 == 0: self.logger.log("Waiting for server to startup, retrying...", level=LogLevel.INFO) if n_retries > 60: raise RuntimeError("Unable to connect to sandbox") time.sleep(1.0) @classmethod def _strip_ansi_colors(cls, text: str) -> str: """Remove ansi colors from text.""" return cls._ANSI_ESCAPE.sub("", text) class BlaxelExecutor(RemotePythonExecutor): """ Remote Python code executor in a Blaxel sandbox. Blaxel provides fast-launching virtual machines that start from hibernation in under 25ms and scale back to zero after inactivity while maintaining memory state. Args: additional_imports (`list[str]`): Additional Python packages to install. logger (`Logger`): Logger to use for output and errors. allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` if you fully trust the execution environment and need backward compatibility with custom types. sandbox_name (`str`, *optional*): Name for the sandbox. Defaults to "smolagent-executor". image (`str`, default `"blaxel/jupyter-notebook"`): Docker image to use. memory (`int`, default `4096`): Memory allocation in MB. ttl (`str`, *optional*): Time to live in seconds. region (`str`, *optional*): Deployment region. If not specified, Blaxel chooses default. """ def __init__( self, additional_imports: list[str], logger, allow_pickle: bool = False, sandbox_name: str | None = None, image: str = "blaxel/jupyter-notebook", memory: int = 4096, ttl: str | None = None, region: Optional[str] = None, ): super().__init__(additional_imports, logger, allow_pickle=allow_pickle) try: import blaxel # noqa: F401 except ModuleNotFoundError: raise ModuleNotFoundError( "Please install 'blaxel' extra to use BlaxelExecutor: `pip install 'smolagents[blaxel]'`" ) self.sandbox_name = sandbox_name or f"smolagent-executor-{uuid.uuid4().hex[:8]}" self.image = image self.memory = memory self.region = region self.port = 8888 self._cleaned_up = False # Flag to prevent double cleanup # Prepare sandbox creation parameters token = secrets.token_urlsafe(16) sandbox_config = { "metadata": { "name": self.sandbox_name, }, "spec": { "runtime": {"image": image, "memory": memory, "ports": [{"target": self.port}]}, }, } if region: sandbox_config["spec"]["region"] = region if ttl: sandbox_config["spec"]["runtime"]["ttl"] = ttl # Create the sandbox try: # Create sandbox environment on Blaxel self.sandbox = BlaxelExecutor._create_sandbox(sandbox_config) # Create kernel via HTTP from blaxel.core import settings kernel_id = _create_kernel_http( f"{self.sandbox.metadata.url}/port/{self.port}/api/kernels?token={token}", self.logger, headers=settings.headers, ) # Set up websocket URL # Convert http/https to ws/wss ws_scheme = "wss" if self.sandbox.metadata.url.startswith("https") else "ws" ws_base = self.sandbox.metadata.url.replace("https://", "").replace("http://", "") self.ws_url = f"{ws_scheme}://{ws_base}/port/{self.port}/api/kernels/{kernel_id}/channels?token={token}" # Install additional packages self.installed_packages = self.install_packages(additional_imports) self.logger.log("Blaxel is running", level=LogLevel.INFO) except Exception as e: self.cleanup() raise RuntimeError(f"Failed to initialize Blaxel sandbox: {e}") from e @staticmethod def _create_sandbox(config): """Helper method to create sandbox asynchronously.""" from blaxel.core import SandboxInstance from blaxel.core.client import client from blaxel.core.client.api.compute import create_sandbox response = create_sandbox.sync(client=client, body=config) return SandboxInstance(response) def run_code_raise_errors(self, code: str) -> CodeOutput: """ Execute Python code in the Blaxel sandbox and return the result. Args: code (`str`): Python code to execute. Returns: `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. """ from blaxel.core import settings from websocket import create_connection headers = [] for key, value in settings.headers.items(): headers.append(f"{key}: {value}") with closing(create_connection(self.ws_url, header=headers)) as ws: return _websocket_run_code_raise_errors(code, ws, self.logger, self.allow_pickle) def install_packages(self, additional_imports: list[str]) -> list[str]: """Helper method to install packages asynchronously.""" if not additional_imports: return [] from blaxel.core import settings from blaxel.core.sandbox.client import client from blaxel.core.sandbox.client.api.process import get_process_identifier, post_process from blaxel.core.sandbox.client.models import ErrorResponse, ProcessResponse try: client.with_base_url(self.sandbox.metadata.url) client.with_headers(settings.headers) # Install packages using pip via run_code self.logger.log(f"Installing packages: {', '.join(additional_imports)}", level=LogLevel.INFO) pip_install_code = f"pip install --root-user-action=ignore {' '.join(additional_imports)}" identifier = "install-packages" body = { "name": identifier, "command": pip_install_code, } post_process.sync(client=client, body=body) status = "running" interval = 1000 max_wait = 600000 start_time = time.time() * 1000 logs = "" exit_code = 0 while status == "running": if (time.time() * 1000) - start_time > max_wait: raise Exception("Process did not finish in time") data = get_process_identifier.sync(identifier, client=client) if isinstance(data, ProcessResponse): status = data.status or "running" exit_code = data.exit_code logs = data.logs elif isinstance(data, ErrorResponse): raise Exception(f"Failed to install packages: {data.message}") else: raise Exception(f"Unknown response: {data}") if status == "running": time.sleep(interval / 1000) # Convert to seconds if exit_code != 0: self.logger.log_error(f"Failed to install packages (exit code {exit_code}): {logs}") return [] self.logger.log(f"Successfully installed packages: {', '.join(additional_imports)}", level=LogLevel.INFO) return additional_imports except Exception as e: self.logger.log_error(f"Error installing packages: {e}") return [] def _delete_sandbox(self): """Delete sandbox using Blaxel's sync API and wait for completion.""" from blaxel.core.client import client from blaxel.core.client.api.compute import delete_sandbox self.logger.log(f"Requesting sandbox {self.sandbox_name} deletion...", level=LogLevel.INFO) delete_sandbox.sync(client=client, sandbox_name=self.sandbox_name) def cleanup(self): """Sync wrapper to clean up sandbox and resources.""" # Prevent double cleanup if self._cleaned_up: return self.logger.log("Shutting down sandbox...", level=LogLevel.INFO) self._cleaned_up = True try: self._delete_sandbox() except Exception as e: # Log cleanup errors but don't raise - cleanup should be best-effort self.logger.log(f"Error during cleanup: : {e}", level=LogLevel.INFO) finally: # Always clean up local references if hasattr(self, "sandbox"): del self.sandbox self.logger.log("Sandbox cleanup completed", level=LogLevel.INFO) def delete(self): """Ensure cleanup on deletion.""" self.cleanup() def __del__(self): """Ensure cleanup on deletion.""" try: self.cleanup() except Exception: pass # Silently ignore errors during cleanup class WasmExecutor(RemotePythonExecutor): """ Remote Python code executor in a sandboxed WebAssembly environment powered by Pyodide and Deno. This executor combines Deno's secure runtime with Pyodide's WebAssembly‑compiled Python interpreter to deliver s trong isolation guarantees while enabling full Python execution. Args: additional_imports (`list[str]`): Additional Python packages to install in the Pyodide environment. logger (`Logger`): Logger to use for output and errors. allow_pickle (`bool`, default `False`): Whether to allow pickle serialization for objects that cannot be safely serialized to JSON. - `False` (default, recommended): Only safe JSON serialization is used. Raises error if object cannot be safely serialized. - `True` (legacy mode): Tries safe JSON serialization first, falls back to pickle with warning if needed. **Security Warning:** Pickle deserialization can execute arbitrary code. Only set `allow_pickle=True` if you fully trust the execution environment and need backward compatibility with custom types. deno_path (`str`, default `"deno"`): Path to the Deno executable. If not provided, will use "deno" from PATH. deno_permissions (`list[str]`, *optional*): List of permissions to grant to the Deno runtime. Default is minimal permissions needed for execution. timeout (`int`, default `60`): Timeout in seconds for code execution """ DEFAULT_SERVER_HOST = "127.0.0.1" DEFAULT_SERVER_PORT = 8000 def __init__( self, additional_imports: list[str], logger, allow_pickle: bool = False, deno_path: str = "deno", deno_permissions: list[str] | None = None, timeout: int = 60, ): super().__init__(additional_imports, logger, allow_pickle=allow_pickle) # Check if Deno is installed try: subprocess.run([deno_path, "--version"], capture_output=True, check=True) except (subprocess.SubprocessError, FileNotFoundError): raise RuntimeError( "Deno is not installed or not found in PATH. Please install Deno from https://deno.land/" ) self.deno_path = deno_path self.timeout = timeout self.token = secrets.token_urlsafe(16) self.session = requests.Session() self.session.headers["Authorization"] = f"Bearer {self.token}" self.server_host = self.DEFAULT_SERVER_HOST self.server_port = self.DEFAULT_SERVER_PORT # Create the Deno JavaScript runner file self._create_deno_runner() # Default minimal permissions needed if deno_permissions is None: deno_permissions = [ "allow-net=" + ",".join( [ f"{self.server_host}:{self.server_port}", # allow requests to the local server "cdn.jsdelivr.net:443", # allow loading pyodide packages "pypi.org:443,files.pythonhosted.org:443", # allow pyodide install packages from PyPI ] ), # FS permissions are always scoped to deno_cache_dir (the per-instance temp dir # that cleanup() removes). This replaces the original global ~/.cache/deno # permissions, bounding any write from attacker code to a short-lived # directory that never affects other Deno processes or persists past teardown. # --allow-read: required for Deno to load npm package assets at runtime # (e.g. pyodide.asm.wasm is read via Deno file APIs). # --allow-write: required for pyodide's loadPackage() to cache downloaded # Python packages (e.g. micropip) to the Deno-backed FS. f"allow-read={self.deno_cache_dir}", f"allow-write={self.deno_cache_dir}", ] self.deno_permissions = [f"--{perm}" for perm in deno_permissions] # Start the Deno server self._start_deno_server() # Install additional packages self.installed_packages = self.install_packages(additional_imports) self.logger.log("WasmExecutor is running", level=LogLevel.INFO) def _create_deno_runner(self): """Create the Deno JavaScript file that will run Pyodide and execute Python code.""" # Create an isolated per-executor runtime directory to avoid sharing mutable Deno state self.runner_dir = tempfile.mkdtemp(prefix="pyodide_deno_") self.runner_path = os.path.join(self.runner_dir, "pyodide_runner.js") # Create the JavaScript runner file with open(self.runner_path, "w") as f: f.write(self._build_js_code()) # Isolate Deno's module cache inside the per-instance temp directory so it # cannot affect other Deno processes and is removed when cleanup() runs. self.deno_cache_dir = os.path.join(self.runner_dir, "deno_cache") def _build_js_code(self) -> str: """Render JavaScript runner with injected auth token and with configured server host and port.""" return ( self.JS_CODE_TEMPLATE.replace("__AUTH_TOKEN__", self.token) .replace("__SERVER_HOST__", self.server_host) .replace("__SERVER_PORT__", str(self.server_port)) ) def _start_deno_server(self): """Start the Deno server that will run our JavaScript code.""" cmd = [self.deno_path, "run"] + self.deno_permissions + [self.runner_path] # Start the server process self.server_process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env={**os.environ, "DENO_DIR": self.deno_cache_dir}, ) # Wait for the server to start time.sleep(2) # Give the server time to start # Check if the server started successfully if self.server_process.poll() is not None: stderr = self.server_process.stderr.read() raise RuntimeError(f"Failed to start Deno server: {stderr}") self.server_url = f"http://{self.server_host}:{self.server_port}" # Test the connection try: response = self.session.get(self.server_url) if response.status_code != 200: raise RuntimeError(f"Server responded with status code {response.status_code}: {response.text}") except requests.RequestException as e: raise RuntimeError(f"Failed to connect to Deno server: {e}") def run_code_raise_errors(self, code: str) -> CodeOutput: """ Execute Python code in the Pyodide environment and return the result. Args: code (`str`): Python code to execute. Returns: `CodeOutput`: Code output containing the result, logs, and whether it is the final answer. """ try: # Prepare the request payload payload = { "code": code, "packages": self.installed_packages, } # Send the request to the Deno server response = self.session.post(self.server_url, json=payload, timeout=self.timeout) if response.status_code != 200: raise AgentError(f"Server error: {response.text}", self.logger) result = None is_final_answer = False # Parse the response result_data = response.json() # Process the result if result_data.get("result"): result = result_data.get("result") # Check for execution errors elif result_data.get("error"): error = result_data["error"] if ( error.get("pythonExceptionType") == RemotePythonExecutor.FINAL_ANSWER_EXCEPTION and "pythonExceptionValue" in error ): result = self._deserialize_final_answer(error["pythonExceptionValue"], self.allow_pickle) is_final_answer = True else: error_message = f"{error.get('name', 'Error')}: {error.get('message', 'Unknown error')}" if "stack" in error: error_message += f"\n{error['stack']}" raise AgentError(error_message, self.logger) # Get the execution logs execution_logs = result_data.get("stdout", "") # Handle image results if isinstance(result, dict) and result.get("type") == "image": image_data = result.get("data", "") decoded_bytes = base64.b64decode(image_data.encode("utf-8")) return PIL.Image.open(BytesIO(decoded_bytes)), execution_logs return CodeOutput(output=result, logs=execution_logs, is_final_answer=is_final_answer) except requests.RequestException as e: raise AgentError(f"Failed to communicate with Deno server: {e}", self.logger) def install_packages(self, additional_imports: list[str]) -> list[str]: """ Install additional Python packages in the Pyodide environment. Args: additional_imports (`list[str]`): Package names to install. Returns: list[str]: Installed packages. """ # In Pyodide, we don't actually install packages here, but we keep track of them # to load them when executing code # TODO: Install here instead? self.logger.log(f"Adding packages to load: {', '.join(additional_imports)}", level=LogLevel.INFO) return additional_imports def cleanup(self): """Clean up resources used by the executor.""" if hasattr(self, "session"): self.session.close() if hasattr(self, "server_process") and self.server_process: self.logger.log("Stopping Deno server...", level=LogLevel.INFO) self.server_process.terminate() try: self.server_process.wait(timeout=5) except subprocess.TimeoutExpired: self.server_process.kill() # Remove the temporary directory if hasattr(self, "runner_dir") and os.path.exists(self.runner_dir): import shutil shutil.rmtree(self.runner_dir) def delete(self): """Ensure cleanup on deletion.""" self.cleanup() JS_CODE_TEMPLATE = dedent("""\ // pyodide_runner.js - Runs Python code in Pyodide within Deno import { serve } from "https://deno.land/std/http/server.ts"; import { loadPyodide } from "npm:pyodide"; const AUTH_TOKEN = "__AUTH_TOKEN__"; // Initialize Pyodide instance const pyodidePromise = loadPyodide(); // Function to execute Python code and return the result async function executePythonCode(code) { const pyodide = await pyodidePromise; // Create a capture for stdout pyodide.runPython(` import sys import io sys.stdout = io.StringIO() `); // Execute the code and capture any errors let result = null; let error = null; let stdout = ""; try { // Execute the code result = await pyodide.runPythonAsync(code); // Get captured stdout stdout = pyodide.runPython("sys.stdout.getvalue()"); } catch (e) { error = { name: e.constructor.name, message: e.message, stack: e.stack }; // Extract Python exception details if (e.constructor.name === "PythonError") { // Get the Python exception type from the error message: at the end of the traceback const errorMatch = e.message.match(/\\n([^:]+Exception): /); if (errorMatch) { error.pythonExceptionType = errorMatch[1].split(".").pop(); } // If the error is a FinalAnswerException, extract its the encoded value if (error.pythonExceptionType === "FinalAnswerException") { // Extract the base64 encoded value from the error message const valueMatch = e.message.match(/FinalAnswerException: (.*?)(?:\\n|$)/); if (valueMatch) { error.pythonExceptionValue = valueMatch[1]; } } } } return { result, stdout, error }; } // Start a simple HTTP server to receive code execution requests serve(async (req) => { const authHeader = req.headers.get("Authorization"); if (!authHeader || authHeader !== `Bearer ${AUTH_TOKEN}`) { return new Response("Unauthorized", { status: 401 }); } if (req.method === "POST") { try { const body = await req.json(); const { code, packages = [] } = body; // Load any requested packages if (packages && packages.length > 0) { const pyodide = await pyodidePromise; //await pyodide.loadPackagesFromImports(code); await pyodide.loadPackage("micropip"); const micropip = pyodide.pyimport("micropip"); try { await micropip.install(packages); } catch (e) { console.error(`Failed to load package ${packages}: ${e.message}`); } } const result = await executePythonCode(code); return new Response(JSON.stringify(result), { headers: { "Content-Type": "application/json" } }); } catch (e) { return new Response(JSON.stringify({ error: e.message }), { status: 500, headers: { "Content-Type": "application/json" } }); } } return new Response("Pyodide-Deno Executor is running. Send POST requests with code to execute.", { headers: { "Content-Type": "text/plain" } }); }, { hostname: "__SERVER_HOST__", port: __SERVER_PORT__ }); """)