base_agent/tools/code_executor.py

# ════════════════════════════════════════════════════════════════
# tools/code_executor.py
# ════════════════════════════════════════════════════════════════
"""沙箱代码执行工具（从配置读取 timeout / sandbox）"""

import io
import contextlib
import time
from tools.base_tool import BaseTool, ToolResult
from config.settings import settings


class CodeExecutorTool(BaseTool):
    name        = "code_executor"
    description = "在沙箱环境中执行 Python 代码片段，返回标准输出"
    parameters  = {
        "code":    {"type": "string",  "description": "要执行的 Python 代码"},
        "timeout": {"type": "integer", "description": "超时时间（秒）"},
    }

    _SAFE_BUILTINS = {
        "print": print, "range": range, "len": len,
        "int": int, "float": float, "str": str, "list": list,
        "dict": dict, "tuple": tuple, "set": set, "bool": bool,
        "abs": abs, "max": max, "min": min, "sum": sum,
        "enumerate": enumerate, "zip": zip, "map": map,
        "sorted": sorted, "reversed": reversed,
    }

    def __init__(self):
        super().__init__()
        cfg             = settings.tools['code_executor']
        self._timeout   = cfg['timeout']
        self._sandbox   = cfg['sandbox']
        self.logger.debug(
            f"⚙️  CodeExecutor timeout={self._timeout}s, sandbox={self._sandbox}"
        )

    def execute(self, code: str, timeout: int | None = None, **_) -> ToolResult:
        timeout    = timeout or self._timeout
        stdout_buf = io.StringIO()
        start      = time.perf_counter()

        exec_globals = (
            {"__builtins__": self._SAFE_BUILTINS}
            if self._sandbox
            else {"__builtins__": __builtins__}
        )

        try:
            with contextlib.redirect_stdout(stdout_buf):
                exec(compile(code, "<agent_sandbox>", "exec"), exec_globals)  # noqa: S102
            elapsed = (time.perf_counter() - start) * 1000
            output  = stdout_buf.getvalue() or "(无输出)"
            return ToolResult(
                success=True,
                output=f"执行成功 ({elapsed:.1f}ms) [sandbox={self._sandbox}]:\n{output}",
                metadata={"elapsed_ms": elapsed, "sandbox": self._sandbox},
            )
        except Exception as exc:
            return ToolResult(success=False, output=f"执行错误: {type(exc).__name__}: {exc}")