diff --git a/examples/e2b_repl_example.py b/examples/e2b_repl_example.py new file mode 100644 index 000000000..baa3b079a --- /dev/null +++ b/examples/e2b_repl_example.py @@ -0,0 +1,28 @@ +import os + +from dotenv import load_dotenv + +from rlm import RLM +from rlm.logger import RLMLogger + +load_dotenv() + +logger = RLMLogger(log_dir="./logs") + +rlm = RLM( + backend="openai", + backend_kwargs={ + "api_key": os.getenv("OPENAI_API_KEY"), + "model_name": "gpt-5-nano", + }, + environment="e2b", + environment_kwargs={ + "timeout": 300, + }, + max_depth=1, + logger=logger, + verbose=True, +) + +result = rlm.completion("Using your code, solve 2^(2^(2^(2))). Show your work in Python.") +print(result.response) diff --git a/pyproject.toml b/pyproject.toml index 99aff8a4c..983e14653 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ Issues = "https://github.com/alexzhang13/rlm/issues" [project.optional-dependencies] modal = ["modal>=0.73.0", "dill>=0.3.7"] +e2b = ["e2b-code-interpreter>=0.0.11", "dill>=0.3.7"] daytona = ["daytona>=0.128.1", "dill>=0.3.7"] prime = ["prime-sandboxes>=0.2.0", "dill>=0.3.7"] diff --git a/rlm/core/types.py b/rlm/core/types.py index f20474d40..195bd8be3 100644 --- a/rlm/core/types.py +++ b/rlm/core/types.py @@ -13,7 +13,7 @@ "azure_openai", "gemini", ] -EnvironmentType = Literal["local", "docker", "modal", "prime", "daytona"] +EnvironmentType = Literal["local", "docker", "modal", "prime", "daytona", "e2b"] def _serialize_value(value: Any) -> Any: diff --git a/rlm/environments/__init__.py b/rlm/environments/__init__.py index eed91b302..0b28629e5 100644 --- a/rlm/environments/__init__.py +++ b/rlm/environments/__init__.py @@ -7,12 +7,12 @@ def get_environment( - environment: Literal["local", "modal", "docker", "daytona", "prime"], + environment: Literal["local", "modal", "docker", "daytona", "prime", "e2b"], environment_kwargs: dict[str, Any], ) -> BaseEnv: """ Routes a specific environment and the args (as a dict) to the appropriate environment if supported. - Currently supported environments: ['local', 'modal', 'docker', 'daytona', 'prime'] + Currently supported environments: ['local', 'modal', 'docker', 'daytona', 'prime', 'e2b'] """ if environment == "local": return LocalREPL(**environment_kwargs) @@ -32,7 +32,11 @@ def get_environment( from rlm.environments.prime_repl import PrimeREPL return PrimeREPL(**environment_kwargs) + elif environment == "e2b": + from rlm.environments.e2b_repl import E2BREPL + + return E2BREPL(**environment_kwargs) else: raise ValueError( - f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker', 'daytona', 'prime']" + f"Unknown environment: {environment}. Supported: ['local', 'modal', 'docker', 'daytona', 'prime', 'e2b']" ) diff --git a/rlm/environments/e2b_repl.py b/rlm/environments/e2b_repl.py new file mode 100644 index 000000000..aa33975db --- /dev/null +++ b/rlm/environments/e2b_repl.py @@ -0,0 +1,506 @@ +""" +E2B REPL environment that runs Python code in E2B sandboxes. + +Uses the E2B Code Interpreter SDK (https://e2b.dev/docs) for sandbox management. +Follows the same HTTP broker pattern as ModalREPL for LLM communication. +""" + +import base64 +import json +import textwrap +import threading +import time +from typing import Any + +import requests +from e2b_code_interpreter import Sandbox + +from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched +from rlm.core.types import REPLResult, RLMChatCompletion +from rlm.environments.base_env import IsolatedEnv + +# ============================================================================= +# Broker Server Script (runs inside sandbox, handles LLM request queue) +# ============================================================================= + +_BROKER_SCRIPT = textwrap.dedent( + ''' +import json +import threading +import uuid +from flask import Flask, request, jsonify + +app = Flask(__name__) + +# Request queue: {request_id: {"request": {...}, "response": None, "event": Event}} +pending_requests = {} +lock = threading.Lock() + +@app.route("/health") +def health(): + return jsonify({"status": "ok"}) + +@app.route("/enqueue", methods=["POST"]) +def enqueue(): + """Called by sandbox code to submit an LLM request and wait for response.""" + data = request.json + request_id = str(uuid.uuid4()) + event = threading.Event() + + with lock: + pending_requests[request_id] = { + "request": data, + "response": None, + "event": event, + } + + # Wait for response (with timeout) + event.wait(timeout=300) + + with lock: + entry = pending_requests.pop(request_id, None) + + if entry and entry["response"] is not None: + return jsonify(entry["response"]) + else: + return jsonify({"error": "Request timed out"}), 504 + +@app.route("/pending") +def get_pending(): + """Called by E2BREPL to get pending requests.""" + with lock: + pending = [ + {"id": rid, "request": entry["request"]} + for rid, entry in pending_requests.items() + if entry["response"] is None + ] + return jsonify({"pending": pending}) + +@app.route("/respond", methods=["POST"]) +def respond(): + """Called by E2BREPL to submit a response.""" + data = request.json + request_id = data.get("id") + response = data.get("response") + + with lock: + if request_id in pending_requests: + pending_requests[request_id]["response"] = response + pending_requests[request_id]["event"].set() + return jsonify({"status": "ok"}) + + return jsonify({"error": "Request not found"}), 404 + +if __name__ == "__main__": + app.run(host="0.0.0.0", port=8889, threaded=True) +''' +) + + +# ============================================================================= +# Execution Script (runs inside the sandbox for each code block) +# ============================================================================= + + +def _build_exec_script(code: str, broker_port: int = 8888) -> str: + """ + Build a script that executes code with state persistence. + LLM queries go through the local broker server. + """ + code_b64 = base64.b64encode(code.encode()).decode() + + return textwrap.dedent( + f''' +import sys +import io +import json +import base64 +import traceback +import os +import requests + +try: + import dill +except ImportError: + import pickle as dill + +# ============================================================================= +# LLM Query Functions (via local broker) +# ============================================================================= + +BROKER_URL = "http://127.0.0.1:{broker_port}" + +def llm_query(prompt, model=None): + """Query the LM via the broker.""" + try: + response = requests.post( + f"{{BROKER_URL}}/enqueue", + json={{"type": "single", "prompt": prompt, "model": model}}, + timeout=300, + ) + data = response.json() + if data.get("error"): + return f"Error: {{data['error']}}" + return data.get("response", "Error: No response") + except Exception as e: + return f"Error: LM query failed - {{e}}" + + +def llm_query_batched(prompts, model=None): + """Query the LM with multiple prompts.""" + try: + response = requests.post( + f"{{BROKER_URL}}/enqueue", + json={{"type": "batched", "prompts": prompts, "model": model}}, + timeout=300, + ) + data = response.json() + if data.get("error"): + return [f"Error: {{data['error']}}"] * len(prompts) + return data.get("responses", ["Error: No response"] * len(prompts)) + except Exception as e: + return [f"Error: LM query failed - {{e}}"] * len(prompts) + + +# ============================================================================= +# State Management +# ============================================================================= + +STATE_FILE = "/tmp/rlm_state.dill" + +def load_state(): + if os.path.exists(STATE_FILE): + try: + with open(STATE_FILE, "rb") as f: + return dill.load(f) + except: + pass + return {{}} + +def save_state(state): + clean_state = {{}} + for k, v in state.items(): + if k.startswith("_"): + continue + try: + dill.dumps(v) + clean_state[k] = v + except: + pass + with open(STATE_FILE, "wb") as f: + dill.dump(clean_state, f) + +def serialize_locals(state): + result = {{}} + for k, v in state.items(): + if k.startswith("_"): + continue + try: + result[k] = repr(v) + except: + result[k] = f"<{{type(v).__name__}}>" + return result + +# ============================================================================= +# Execution +# ============================================================================= + +_locals = load_state() + +def FINAL_VAR(variable_name): + variable_name = variable_name.strip().strip("\\"\\'") + if variable_name in _locals: + return str(_locals[variable_name]) + return f"Error: Variable '{{variable_name}}' not found" + +_globals = {{ + "__builtins__": __builtins__, + "__name__": "__main__", + "llm_query": llm_query, + "llm_query_batched": llm_query_batched, + "FINAL_VAR": FINAL_VAR, +}} + +code = base64.b64decode("{code_b64}").decode() + +stdout_buf = io.StringIO() +stderr_buf = io.StringIO() +old_stdout, old_stderr = sys.stdout, sys.stderr + +try: + sys.stdout = stdout_buf + sys.stderr = stderr_buf + combined = {{**_globals, **_locals}} + exec(code, combined, combined) + for key, value in combined.items(): + if key not in _globals and not key.startswith("_"): + _locals[key] = value +except Exception as e: + traceback.print_exc(file=stderr_buf) +finally: + sys.stdout = old_stdout + sys.stderr = old_stderr + +save_state(_locals) + +result = {{ + "stdout": stdout_buf.getvalue(), + "stderr": stderr_buf.getvalue(), + "locals": serialize_locals(_locals), +}} +print(json.dumps(result)) +''' + ) + + +class E2BREPL(IsolatedEnv): + """ + E2B REPL environment that runs Python code in E2B sandboxes. + + Uses E2B's public URL feature for LLM communication: + - Sandbox runs a broker server exposed via get_host() + - E2BREPL polls the broker for pending LLM requests + - E2BREPL forwards requests to the LM handler and posts responses back + """ + + BROKER_PORT = 8889 # Use 8889 since E2B Code Interpreter uses 8888 for Jupyter + + def __init__( + self, + timeout: int = 300, # 5 minutes default (E2B uses seconds) + lm_handler_address: tuple[str, int] | None = None, + context_payload: dict | list | str | None = None, + setup_code: str | None = None, + persistent: bool = False, + **kwargs: Any, + ): + if persistent: + raise NotImplementedError( + "Persistent REPLs are currently not supported for environment: E2BREPL" + ) + super().__init__(persistent=persistent, **kwargs) + + self.timeout = timeout + self.lm_handler_address = lm_handler_address + + # Sandbox state + self.sandbox: Sandbox | None = None + self.sandbox_id: str | None = None + self.broker_url: str | None = None + self.broker_process = None + + # Polling thread for LLM requests + self.poller_thread: threading.Thread | None = None + self.poller_stop = threading.Event() + self.pending_llm_calls: list[RLMChatCompletion] = [] + self._calls_lock = threading.Lock() + + self.setup() + + if context_payload is not None: + self.load_context(context_payload) + + if setup_code: + self.execute_code(setup_code) + + def setup(self): + """Create the E2B sandbox, broker, and start polling.""" + # Create the sandbox using the recommended Sandbox.create() API + self.sandbox = Sandbox.create(timeout=self.timeout) + self.sandbox_id = self.sandbox.sandbox_id + + # Install dependencies for the broker + self.sandbox.commands.run("pip install flask requests dill") + + # Write the broker script to the sandbox + self.sandbox.files.write("/tmp/broker.py", _BROKER_SCRIPT) + + # Start the broker as a background process + self.broker_process = self.sandbox.commands.run( + "python /tmp/broker.py", + background=True, + ) + + # Wait for broker to be ready + self._wait_for_broker() + + # Get the public URL for the broker port + host = self.sandbox.get_host(self.BROKER_PORT) + self.broker_url = f"https://{host}" + + # Start polling thread if we have an LM handler + if self.lm_handler_address and self.broker_url: + self.poller_stop.clear() + self.poller_thread = threading.Thread(target=self._poll_broker, daemon=True) + self.poller_thread.start() + + def _wait_for_broker(self, max_attempts: int = 30): + """Wait for the broker to be ready by checking health endpoint.""" + health_check_cmd = ( + f'python -c "import requests; ' + f"r = requests.get('http://127.0.0.1:{self.BROKER_PORT}/health', timeout=2); " + f'print(r.text)"' + ) + + for _ in range(max_attempts): + time.sleep(1) + try: + result = self.sandbox.commands.run(health_check_cmd) + if result.stdout and "ok" in result.stdout.lower(): + return + except Exception: + pass + + raise RuntimeError("Broker failed to start within timeout") + + def _poll_broker(self): + """Poll the broker for pending LLM requests and handle them.""" + while not self.poller_stop.is_set(): + try: + # Get pending requests + resp = requests.get( + f"{self.broker_url}/pending", + timeout=5, + ) + pending = resp.json().get("pending", []) + + for item in pending: + request_id = item["id"] + req_data = item["request"] + + # Handle the request + response = self._handle_llm_request(req_data) + + # Send response back + requests.post( + f"{self.broker_url}/respond", + json={"id": request_id, "response": response}, + timeout=10, + ) + + except requests.exceptions.RequestException: + pass + except Exception: + pass + + time.sleep(0.1) + + def _handle_llm_request(self, req_data: dict) -> dict: + """Handle an LLM request from the sandbox.""" + req_type = req_data.get("type") + model = req_data.get("model") + + if req_type == "single": + prompt = req_data.get("prompt") + request = LMRequest(prompt=prompt, model=model) + response = send_lm_request(self.lm_handler_address, request) + + if not response.success: + return {"error": response.error} + + # Track the call + with self._calls_lock: + self.pending_llm_calls.append(response.chat_completion) + + return {"response": response.chat_completion.response} + + elif req_type == "batched": + prompts = req_data.get("prompts", []) + responses = send_lm_request_batched(self.lm_handler_address, prompts, model=model) + + results = [] + for resp in responses: + if not resp.success: + results.append(f"Error: {resp.error}") + else: + with self._calls_lock: + self.pending_llm_calls.append(resp.chat_completion) + results.append(resp.chat_completion.response) + + return {"responses": results} + + return {"error": "Unknown request type"} + + def load_context(self, context_payload: dict | list | str): + """Load context into the sandbox environment.""" + if isinstance(context_payload, str): + escaped = context_payload.replace("\\", "\\\\").replace('"""', '\\"\\"\\"') + context_code = f'context = """{escaped}"""' + else: + context_json = json.dumps(context_payload) + escaped_json = context_json.replace("\\", "\\\\").replace("'", "\\'") + context_code = f"import json; context = json.loads('{escaped_json}')" + + self.execute_code(context_code) + + def execute_code(self, code: str) -> REPLResult: + """Execute code in the E2B sandbox and return result.""" + start_time = time.perf_counter() + + # Clear pending LLM calls + with self._calls_lock: + self.pending_llm_calls.clear() + + # Build and write the script to sandbox + script = _build_exec_script(code, self.BROKER_PORT) + self.sandbox.files.write("/tmp/run_script.py", script) + + # Run the script + result = self.sandbox.commands.run("python /tmp/run_script.py", timeout=600) + stdout = result.stdout or "" + stderr = result.stderr or "" + + # Collect LLM calls made during this execution + with self._calls_lock: + pending_calls = self.pending_llm_calls.copy() + self.pending_llm_calls.clear() + + execution_time = time.perf_counter() - start_time + + # Parse the JSON result + try: + lines = stdout.strip().split("\n") + result_json = lines[-1] if lines else "{}" + parsed = json.loads(result_json) + + return REPLResult( + stdout=parsed.get("stdout", ""), + stderr=parsed.get("stderr", "") + stderr, + locals=parsed.get("locals", {}), + execution_time=execution_time, + rlm_calls=pending_calls, + ) + except json.JSONDecodeError: + return REPLResult( + stdout=stdout, + stderr=stderr or "Failed to parse execution result", + locals={}, + execution_time=execution_time, + rlm_calls=pending_calls, + ) + + def cleanup(self): + """Terminate the sandbox and stop polling.""" + # Stop the poller thread + if self.poller_thread is not None: + self.poller_stop.set() + self.poller_thread.join(timeout=2) + self.poller_thread = None + + # Kill the sandbox + if self.sandbox is not None: + try: + self.sandbox.kill() + except Exception: + pass + self.sandbox = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.cleanup() + return False + + def __del__(self): + self.cleanup()