zcbot/tools/run_python.py

109 lines
4.2 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 the .py under scripts/ (e.g. scripts/analyze.py) first, "
"then execute it so the full source stays in files instead of conversation history.\n"
"Use inline code only for short throwaway snippets (a quick calc / one-line probe) — these run from a temp file and leave no trace.\n"
"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 (task_dir); relative paths resolve against it. "
"Keep process scripts in scripts/; write deliverables to task_dir root or the SKILL-specified path.\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 (relative to task_dir). Prefer this for non-trivial code; keep such scripts under scripts/.",
},
"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))