LZH-YS1998
131 lines
5.1 KiB
Python
131 lines
5.1 KiB
Python
"""Python code execution tool."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from opc.layer4_tools.execution_context import (
|
|
build_subprocess_env,
|
|
resolve_python_executable,
|
|
resolve_task_execution_context,
|
|
wrap_command_for_context,
|
|
)
|
|
from opc.layer4_tools.output_budget import clip_text, persist_tool_result
|
|
from opc.layer4_tools.registry import ToolDefinition
|
|
|
|
|
|
async def python_exec(
|
|
code: str,
|
|
timeout: int = 60,
|
|
task: Any | None = None,
|
|
on_progress: Any = None,
|
|
) -> dict[str, Any]:
|
|
"""Execute Python code in a subprocess and capture output."""
|
|
_ = on_progress
|
|
context = resolve_task_execution_context(task)
|
|
workspace_root = str(context.get("workspace_root", "") or "").strip()
|
|
temp_dir = workspace_root if workspace_root and Path(workspace_root).exists() else None
|
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, dir=temp_dir) as f:
|
|
f.write(code)
|
|
tmp_path = f.name
|
|
|
|
proc: asyncio.subprocess.Process | None = None
|
|
try:
|
|
executable = resolve_python_executable(context)
|
|
cwd = workspace_root or os.getcwd()
|
|
env = build_subprocess_env(context)
|
|
wrapped_args, sandbox_meta = wrap_command_for_context(
|
|
[executable, tmp_path],
|
|
cwd=str(Path(cwd).resolve()),
|
|
context=context,
|
|
)
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*wrapped_args,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
cwd=str(Path(cwd).resolve()),
|
|
env=env,
|
|
)
|
|
stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
|
stdout_text = stdout.decode("utf-8", errors="replace")
|
|
stderr_text = stderr.decode("utf-8", errors="replace")
|
|
stdout_clip = clip_text(stdout_text, limit=30000, marker="python stdout truncated")
|
|
stderr_clip = clip_text(stderr_text, limit=10000, marker="python stderr truncated")
|
|
stdout_persisted = (
|
|
persist_tool_result(stdout_text, tool_name="python_exec_stdout", task=task, extension="txt")
|
|
if stdout_clip.truncated else {}
|
|
)
|
|
stderr_persisted = (
|
|
persist_tool_result(stderr_text, tool_name="python_exec_stderr", task=task, extension="txt")
|
|
if stderr_clip.truncated else {}
|
|
)
|
|
return {
|
|
"stdout": stdout_clip.text,
|
|
"stderr": stderr_clip.text,
|
|
"stdout_truncated": stdout_clip.truncated,
|
|
"stdout_omitted_chars": stdout_clip.omitted_chars,
|
|
"stderr_truncated": stderr_clip.truncated,
|
|
"stderr_omitted_chars": stderr_clip.omitted_chars,
|
|
"full_stdout_path": stdout_persisted.get("full_output_path", ""),
|
|
"full_stderr_path": stderr_persisted.get("full_output_path", ""),
|
|
"exit_code": proc.returncode,
|
|
"sandbox": sandbox_meta,
|
|
"execution_context": {
|
|
"workspace_root": workspace_root,
|
|
"venv_path": str(context.get("venv_path", "") or ""),
|
|
"python_executable": executable,
|
|
"preparation_error": str(context.get("preparation_error", "") or ""),
|
|
"sandbox": dict(context.get("sandbox", {}) or {}),
|
|
},
|
|
}
|
|
except asyncio.TimeoutError:
|
|
if proc is not None:
|
|
proc.kill()
|
|
return {
|
|
"error": f"Execution timed out after {timeout}s",
|
|
"exit_code": -1,
|
|
"execution_context": {
|
|
"workspace_root": workspace_root,
|
|
"venv_path": str(context.get("venv_path", "") or ""),
|
|
"python_executable": resolve_python_executable(context),
|
|
"preparation_error": str(context.get("preparation_error", "") or ""),
|
|
},
|
|
}
|
|
except RuntimeError as exc:
|
|
return {
|
|
"error": str(exc),
|
|
"exit_code": -1,
|
|
"execution_context": {
|
|
"workspace_root": workspace_root,
|
|
"venv_path": str(context.get("venv_path", "") or ""),
|
|
"python_executable": resolve_python_executable(context),
|
|
"preparation_error": str(context.get("preparation_error", "") or ""),
|
|
"sandbox": dict(context.get("sandbox", {}) or {}),
|
|
},
|
|
}
|
|
finally:
|
|
Path(tmp_path).unlink(missing_ok=True)
|
|
|
|
|
|
def create_python_tool() -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name="python_exec",
|
|
description="Execute Python code and return stdout/stderr. Use for data processing, calculations, or testing.",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"code": {"type": "string", "description": "Python code to execute"},
|
|
"timeout": {"type": "integer", "description": "Timeout in seconds", "default": 60},
|
|
},
|
|
"required": ["code"],
|
|
},
|
|
func=python_exec,
|
|
category="compute",
|
|
self_bounded_output=True,
|
|
max_result_chars=80_000,
|
|
)
|