Add python_exec tool: run Python code in subprocess with ctx injection and timeout
Executes code in a fresh python3 subprocess (max 30s, default 10s). Pre-injects a 'ctx' dict with: chat_id, message_id, message_text, args, from_user, replied_message. stdout+stderr returned; output capped at 3000 chars. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,10 +1,12 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
import inspect
|
import inspect
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||||
|
|
||||||
@@ -19,7 +21,7 @@ from app.ai.image_understand import (
|
|||||||
)
|
)
|
||||||
from app.ai.parser import AIInvocation
|
from app.ai.parser import AIInvocation
|
||||||
from app.ai.prompts import PromptBook
|
from app.ai.prompts import PromptBook
|
||||||
from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, mcp_tools_to_openai_schemas
|
from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, PYTHON_EXEC_SCHEMA, mcp_tools_to_openai_schemas
|
||||||
from app.config import AIConfig, AIModelConfig, AIProviderConfig
|
from app.config import AIConfig, AIModelConfig, AIProviderConfig
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -327,11 +329,92 @@ class AIService:
|
|||||||
|
|
||||||
async def _build_tool_list(self) -> list[dict]:
|
async def _build_tool_list(self) -> list[dict]:
|
||||||
mcp_schemas = await self._fetch_mcp_tool_schemas()
|
mcp_schemas = await self._fetch_mcp_tool_schemas()
|
||||||
return FILE_TOOL_SCHEMAS + mcp_schemas
|
return FILE_TOOL_SCHEMAS + [PYTHON_EXEC_SCHEMA] + mcp_schemas
|
||||||
|
|
||||||
def _get_allowed_roots(self) -> tuple[Path, ...]:
|
def _get_allowed_roots(self) -> tuple[Path, ...]:
|
||||||
return (self.config.workspace_dir, self.config.prompt.root)
|
return (self.config.workspace_dir, self.config.prompt.root)
|
||||||
|
|
||||||
|
def _build_exec_context(self, ctx: "CommandContext") -> dict[str, Any]:
|
||||||
|
actor = ctx.actor
|
||||||
|
replied = ctx.replied_message
|
||||||
|
replied_user = ctx.replied_user
|
||||||
|
return {
|
||||||
|
"chat_id": getattr(getattr(ctx.message, "chat", None), "id", None),
|
||||||
|
"message_id": getattr(ctx.message, "id", None),
|
||||||
|
"message_text": getattr(ctx.message, "text", None) or "",
|
||||||
|
"args": ctx.args,
|
||||||
|
"from_user": {
|
||||||
|
"id": getattr(actor, "id", None),
|
||||||
|
"first_name": getattr(actor, "first_name", None),
|
||||||
|
"last_name": getattr(actor, "last_name", None),
|
||||||
|
"username": getattr(actor, "username", None),
|
||||||
|
} if actor else None,
|
||||||
|
"replied_message": {
|
||||||
|
"id": getattr(replied, "id", None),
|
||||||
|
"text": getattr(replied, "text", None) or getattr(replied, "caption", None) or "",
|
||||||
|
"from_user": {
|
||||||
|
"id": getattr(replied_user, "id", None),
|
||||||
|
"first_name": getattr(replied_user, "first_name", None),
|
||||||
|
"username": getattr(replied_user, "username", None),
|
||||||
|
} if replied_user else None,
|
||||||
|
} if replied else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
async def _execute_python_exec(
|
||||||
|
self,
|
||||||
|
arguments: dict,
|
||||||
|
ctx: "CommandContext",
|
||||||
|
) -> str:
|
||||||
|
code = str(arguments.get("code", "")).strip()
|
||||||
|
if not code:
|
||||||
|
raise ValueError("python_exec: 'code' is required")
|
||||||
|
|
||||||
|
raw_timeout = arguments.get("timeout", 10)
|
||||||
|
try:
|
||||||
|
timeout = float(raw_timeout)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
timeout = 10.0
|
||||||
|
timeout = min(max(timeout, 1.0), 30.0)
|
||||||
|
|
||||||
|
context = self._build_exec_context(ctx)
|
||||||
|
context_json = json.dumps(context, ensure_ascii=False, default=str)
|
||||||
|
context_line = f"import json as _json; ctx = _json.loads({context_json!r})\n"
|
||||||
|
full_code = (context_line + code).encode("utf-8")
|
||||||
|
|
||||||
|
try:
|
||||||
|
proc = await asyncio.create_subprocess_exec(
|
||||||
|
sys.executable, "-",
|
||||||
|
stdin=asyncio.subprocess.PIPE,
|
||||||
|
stdout=asyncio.subprocess.PIPE,
|
||||||
|
stderr=asyncio.subprocess.PIPE,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
stdout, stderr = await asyncio.wait_for(
|
||||||
|
proc.communicate(input=full_code), timeout=timeout
|
||||||
|
)
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
proc.kill()
|
||||||
|
await proc.wait()
|
||||||
|
return f"Execution timed out after {timeout:.0f}s"
|
||||||
|
except OSError as exc:
|
||||||
|
raise ValueError(f"python_exec: failed to start interpreter: {exc}") from exc
|
||||||
|
|
||||||
|
out = stdout.decode("utf-8", errors="replace").rstrip()
|
||||||
|
err = stderr.decode("utf-8", errors="replace").rstrip()
|
||||||
|
|
||||||
|
parts: list[str] = []
|
||||||
|
if out:
|
||||||
|
parts.append(out)
|
||||||
|
if err:
|
||||||
|
parts.append(f"[stderr]\n{err}")
|
||||||
|
if not parts:
|
||||||
|
parts.append(f"[exit {proc.returncode}]" if proc.returncode else "(no output)")
|
||||||
|
|
||||||
|
result = "\n".join(parts)
|
||||||
|
if len(result) > 3000:
|
||||||
|
result = result[:3000] + "\n…(truncated)"
|
||||||
|
return result
|
||||||
|
|
||||||
async def _execute_send_file(
|
async def _execute_send_file(
|
||||||
self,
|
self,
|
||||||
arguments: dict,
|
arguments: dict,
|
||||||
@@ -373,7 +456,9 @@ class AIService:
|
|||||||
return f"Tool {name!r}: invalid JSON arguments: {exc}"
|
return f"Tool {name!r}: invalid JSON arguments: {exc}"
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if name == "send_file":
|
if name == "python_exec":
|
||||||
|
result = await self._execute_python_exec(arguments, ctx)
|
||||||
|
elif name == "send_file":
|
||||||
result = await self._execute_send_file(arguments, ctx, allowed_roots)
|
result = await self._execute_send_file(arguments, ctx, allowed_roots)
|
||||||
elif name in FILE_TOOL_DISPATCH:
|
elif name in FILE_TOOL_DISPATCH:
|
||||||
result = await FILE_TOOL_DISPATCH[name](arguments, allowed_roots)
|
result = await FILE_TOOL_DISPATCH[name](arguments, allowed_roots)
|
||||||
|
|||||||
@@ -184,6 +184,33 @@ def _walk(
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
PYTHON_EXEC_SCHEMA: dict[str, Any] = {
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "python_exec",
|
||||||
|
"description": (
|
||||||
|
"Execute Python code and return stdout/stderr output. "
|
||||||
|
"A 'ctx' dict is pre-injected with message context: "
|
||||||
|
"chat_id, message_id, message_text, args, from_user, replied_message. "
|
||||||
|
"Use print() to produce output."
|
||||||
|
),
|
||||||
|
"parameters": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"code": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Python code to execute",
|
||||||
|
},
|
||||||
|
"timeout": {
|
||||||
|
"type": "number",
|
||||||
|
"description": "Execution timeout in seconds (default 10, max 30)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["code"],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
|
ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
|
||||||
|
|
||||||
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {
|
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {
|
||||||
|
|||||||
Reference in New Issue
Block a user