zcbot/tools/run_python.py

107 lines
3.9 KiB
Python

"""run_python: 在 subprocess 里执行 Python 代码 (Hybrid 范式的关键)。
JSON tool call 处理离散操作 (read/edit/shell),run_python 处理连续逻辑
(数据计算、批处理、生成 .pptx/.docx)。让模型自己挑工具。
阶段 1 (本地): subprocess + 工作目录限制 + 敏感环境变量过滤
阶段 2 (后期): Docker / E2B 替换执行后端
"""
from __future__ import annotations
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from .base import Tool, compact_tool_output
_SENSITIVE_PATTERNS = ("API_KEY", "TOKEN", "SECRET", "PASSWORD", "PRIVATE_KEY")
class RunPythonTool(Tool):
name = "run_python"
description = (
"Execute Python code in a subprocess. Returns stdout/stderr/exit_code.\n"
"Use script_path for non-trivial code: write a .py file first, then execute it "
"so the full source stays in files instead of conversation history.\n"
"Use inline code only for short snippets. Good for: data analysis, batch file ops, document generation (.pptx/.docx), "
"matplotlib charts, or any task where Python is more natural than chaining tools.\n"
"Working directory is the agent's base dir. Files you create persist there.\n"
"Available libs (install with shell pip if missing): "
"pandas, numpy, matplotlib, python-pptx, python-docx, requests, pypdf."
)
parameters = {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Short Python source. For longer code, write a .py file and pass script_path.",
},
"script_path": {
"type": "string",
"description": "Path to an existing .py file to execute. Prefer this for non-trivial code.",
},
"timeout": {"type": "integer", "default": 120, "description": "Seconds before kill"},
},
"required": [],
}
def execute(
self,
code: str | None = None,
script_path: str | None = None,
timeout: int = 120,
) -> str:
cleanup_script = False
if script_path:
script = self._resolve(script_path)
if not script.is_file():
return f"[Error] script_path not found: {self._display(script)}"
elif isinstance(code, str):
# 写到临时文件,避免 -c 转义问题
with tempfile.NamedTemporaryFile(
suffix=".py", mode="w", delete=False, encoding="utf-8"
) as f:
f.write(code)
script = Path(f.name)
cleanup_script = True
else:
return "[Error] run_python requires code or script_path"
try:
env = os.environ.copy()
for k in list(env):
u = k.upper()
if any(p in u for p in _SENSITIVE_PATTERNS):
del env[k]
env["PYTHONIOENCODING"] = "utf-8"
env["PYTHONPATH"] = str(self.base_dir) + os.pathsep + env.get("PYTHONPATH", "")
result = subprocess.run(
[sys.executable, str(script)],
cwd=str(self.base_dir),
capture_output=True,
timeout=timeout,
text=True,
encoding="utf-8",
errors="replace",
env=env,
)
except subprocess.TimeoutExpired:
return f"[Error] python script timed out after {timeout}s"
finally:
if cleanup_script:
try:
script.unlink()
except OSError:
pass
parts = []
if result.stdout:
parts.append(f"[stdout]\n{result.stdout.rstrip()}")
if result.stderr:
parts.append(f"[stderr]\n{result.stderr.rstrip()}")
parts.append(f"[exit {result.returncode}]")
return compact_tool_output("\n".join(parts))