feat: 增加python 代码调用,修复一些已知问题

This commit is contained in:
ViperEkura 2026-03-25 10:34:27 +08:00
parent 4499c72ed8
commit ba8b21dd03
4 changed files with 190 additions and 3 deletions

View File

@ -4,7 +4,7 @@ from datetime import datetime
from flask import Blueprint, request
from backend import db
from backend.models import Conversation, Message
from backend.utils.helpers import ok, err, to_dict, get_or_create_default_user
from backend.utils.helpers import ok, err, to_dict, message_to_dict, get_or_create_default_user
from backend.services.chat import ChatService
@ -36,7 +36,7 @@ def message_list(conv_id):
db.session.query(Message.created_at).filter_by(id=cursor).scalar() or datetime.utcnow))
rows = q.order_by(Message.created_at.asc()).limit(limit + 1).all()
items = [to_dict(r) for r in rows[:limit]]
items = [message_to_dict(r) for r in rows[:limit]]
return ok({
"items": items,
"next_cursor": items[-1]["id"] if len(rows) > limit else None,

View File

@ -26,7 +26,7 @@ def init_tools() -> None:
Importing builtin module automatically registers all decorator-defined tools
"""
from backend.tools.builtin import crawler, data, weather, file_ops # noqa: F401
from backend.tools.builtin import code, crawler, data, weather, file_ops # noqa: F401
# Public API exports

View File

@ -0,0 +1,163 @@
"""Safe code execution tool with sandboxing"""
import ast
import subprocess
import sys
import tempfile
import textwrap
from pathlib import Path
from backend.tools.factory import tool
# Whitelist of allowed modules
ALLOWED_MODULES = {
# Standard library - math and data processing
"math", "random", "statistics", "itertools", "functools", "operator",
"collections", "decimal", "fractions", "numbers",
# String processing
"string", "re", "textwrap", "unicodedata",
# Data formats
"json", "csv", "datetime", "time",
# Data structures
"heapq", "bisect", "array", "copy",
# Type related
"typing", "types", "dataclasses",
}
# Blacklist of dangerous builtins
BLOCKED_BUILTINS = {
"eval", "exec", "compile", "open", "input",
"__import__", "globals", "locals", "vars",
"breakpoint", "exit", "quit",
"memoryview", "bytearray",
}
@tool(
name="execute_python",
description="Execute Python code in a sandboxed environment. Supports math, data processing, and string operations. Max execution time: 10 seconds.",
parameters={
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Python code to execute. Only standard library modules allowed."
}
},
"required": ["code"]
},
category="code"
)
def execute_python(arguments: dict) -> dict:
"""
Execute Python code safely with sandboxing.
Security measures:
1. Restricted imports (whitelist)
2. Blocked dangerous builtins
3. Timeout limit (10s)
4. No file system access
5. No network access
"""
code = arguments["code"]
# Security check: detect dangerous imports
dangerous_imports = _check_dangerous_imports(code)
if dangerous_imports:
return {
"success": False,
"error": f"Blocked imports: {', '.join(dangerous_imports)}. Only standard library modules allowed: {', '.join(sorted(ALLOWED_MODULES))}"
}
# Security check: detect dangerous function calls
dangerous_calls = _check_dangerous_calls(code)
if dangerous_calls:
return {
"success": False,
"error": f"Blocked functions: {', '.join(dangerous_calls)}"
}
# Execute in isolated subprocess
try:
result = subprocess.run(
[sys.executable, "-c", _build_safe_code(code)],
capture_output=True,
timeout=10,
cwd=tempfile.gettempdir(),
encoding="utf-8",
env={ # Clear environment variables
"PYTHONIOENCODING": "utf-8",
}
)
if result.returncode == 0:
return {"success": True, "output": result.stdout}
else:
return {"success": False, "error": result.stderr or "Execution failed"}
except subprocess.TimeoutExpired:
return {"success": False, "error": "Execution timeout (10s limit)"}
except Exception as e:
return {"success": False, "error": f"Execution error: {str(e)}"}
def _build_safe_code(code: str) -> str:
"""Build sandboxed code with restricted globals"""
template = textwrap.dedent('''
import builtins
# Block dangerous builtins
_BLOCKED = %r
_safe_builtins = {k: getattr(builtins, k) for k in dir(builtins) if k not in _BLOCKED}
# Create safe namespace
_safe_globals = {
"__builtins__": _safe_builtins,
"__name__": "__main__",
}
# Execute code
exec(%r, _safe_globals)
''').strip()
return template % (BLOCKED_BUILTINS, code)
def _check_dangerous_imports(code: str) -> list:
"""Check for disallowed imports"""
try:
tree = ast.parse(code)
except SyntaxError:
return []
dangerous = []
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
module = alias.name.split(".")[0]
if module not in ALLOWED_MODULES:
dangerous.append(module)
elif isinstance(node, ast.ImportFrom):
if node.module:
module = node.module.split(".")[0]
if module not in ALLOWED_MODULES:
dangerous.append(module)
return dangerous
def _check_dangerous_calls(code: str) -> list:
"""Check for blocked function calls"""
try:
tree = ast.parse(code)
except SyntaxError:
return []
dangerous = []
for node in ast.walk(tree):
if isinstance(node, ast.Call):
if isinstance(node.func, ast.Name):
if node.func.id in BLOCKED_BUILTINS:
dangerous.append(node.func.id)
return dangerous

View File

@ -46,6 +46,30 @@ def to_dict(inst, **extra):
return d
def message_to_dict(msg: Message) -> dict:
"""Convert message to dict with tool calls"""
result = to_dict(msg, thinking_content=msg.thinking_content or None)
# Add tool calls if any
tool_calls = msg.tool_calls.all() if msg.tool_calls else []
if tool_calls:
result["tool_calls"] = [
{
"id": tc.call_id,
"type": "function",
"function": {
"name": tc.tool_name,
"arguments": tc.arguments,
},
"result": tc.result,
"execution_time": tc.execution_time,
}
for tc in tool_calls
]
return result
def record_token_usage(user_id, model, prompt_tokens, completion_tokens):
"""Record token usage"""
today = date.today()