atlas-iac/dockerfiles/hermes-python-sandbox-tool.py

"""Hermes tool for running Python in a separate, credential-free pod."""

from __future__ import annotations

import json
import os
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from tools.registry import registry, tool_error


def _sandbox_url() -> str:
    """Return the tenant-specific sandbox endpoint injected at startup."""
    return os.environ.get("HERMES_CODE_SANDBOX_URL", "").strip()


def _sandbox_available() -> bool:
    """Expose the tool only when this tenant has an isolated endpoint."""
    return _sandbox_url().startswith("http://hermes-chat-sandbox-")


def execute_python_sandbox(code: str) -> str:
    """Execute Python remotely and return the sandbox's structured result."""
    if not isinstance(code, str) or not code.strip():
        return tool_error("Python code is required.")
    if len(code.encode("utf-8")) > 128 * 1024:
        return tool_error("Python code exceeds the 128 KiB request limit.")

    request = Request(
        _sandbox_url(),
        data=json.dumps({"code": code}).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urlopen(request, timeout=55) as response:
            payload = response.read(256 * 1024)
        return payload.decode("utf-8", errors="replace")
    except HTTPError as exc:
        detail = exc.read(4096).decode("utf-8", errors="replace")
        return tool_error(f"Python sandbox rejected the request: {detail or exc.code}")
    except (URLError, TimeoutError, OSError) as exc:
        return tool_error(f"Python sandbox is unavailable: {exc}")


registry.register(
    name="python_sandbox",
    toolset="python_sandbox",
    schema={
        "name": "python_sandbox",
        "description": (
            "Run Python in this user's credential-free computation sandbox. "
            "Its /workspace and /opt/data/workspace paths expose the same private "
            "workspace, and /tmp is writable for temporary verifiers. "
            "Use it for file verification, statistics, probability, Monte Carlo "
            "simulation, data transforms, and calculations. The sandbox has no "
            "Kubernetes, Vault, model-provider credentials, or access to other "
            "users. Public research belongs in web_search."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "A self-contained Python 3 program that prints its result.",
                }
            },
            "required": ["code"],
        },
    },
    handler=lambda args, **_: execute_python_sandbox(args.get("code", "")),
    check_fn=_sandbox_available,
    emoji="🧮",
    max_result_size_chars=200_000,
)