60 lines
2.4 KiB
Python
60 lines
2.4 KiB
Python
|
"""代码执行工具"""
|
|||
|
|
|
|||
|
|
# ════════════════════════════════════════════════════════════════
|
|||
|
|
# tools/code_executor.py — 代码执行工具
|
|||
|
|
# ════════════════════════════════════════════════════════════════
|
|||
|
|
"""
|
|||
|
|
tools/code_executor.py
|
|||
|
|
沙箱代码执行工具：在受限环境中运行 Python 代码片段
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import io
|
|||
|
|
import contextlib
|
|||
|
|
import time
|
|||
|
|
from tools.base_tool import BaseTool, ToolResult
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CodeExecutorTool(BaseTool):
|
|||
|
|
name = "code_executor"
|
|||
|
|
description = "在沙箱环境中执行 Python 代码片段，返回标准输出"
|
|||
|
|
parameters = {
|
|||
|
|
"code": {
|
|||
|
|
"type": "string",
|
|||
|
|
"description": "要执行的 Python 代码",
|
|||
|
|
},
|
|||
|
|
"timeout": {
|
|||
|
|
"type": "integer",
|
|||
|
|
"description": "超时时间（秒），默认 5",
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 沙箱：仅允许安全的内置函数
|
|||
|
|
_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 execute(self, code: str, timeout: int = 5, **_) -> ToolResult:
|
|||
|
|
stdout_buf = io.StringIO()
|
|||
|
|
start_time = time.perf_counter()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 重定向 stdout，捕获 print 输出
|
|||
|
|
with contextlib.redirect_stdout(stdout_buf):
|
|||
|
|
exec( # noqa: S102
|
|||
|
|
compile(code, "<agent_sandbox>", "exec"),
|
|||
|
|
{"__builtins__": self._SAFE_BUILTINS},
|
|||
|
|
)
|
|||
|
|
elapsed = (time.perf_counter() - start_time) * 1000
|
|||
|
|
output = stdout_buf.getvalue() or "(无输出)"
|
|||
|
|
return ToolResult(
|
|||
|
|
success=True,
|
|||
|
|
output=f"执行成功 ({elapsed:.1f}ms):\n{output}",
|
|||
|
|
metadata={"elapsed_ms": elapsed},
|
|||
|
|
)
|
|||
|
|
except Exception as exc:
|
|||
|
|
return ToolResult(success=False, output=f"执行错误: {type(exc).__name__}: {exc}")
|