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 flask import Blueprint, request
from backend import db from backend import db
from backend.models import Conversation, Message 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 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)) 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() 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({ return ok({
"items": items, "items": items,
"next_cursor": items[-1]["id"] if len(rows) > limit else None, "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 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 # 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 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): def record_token_usage(user_id, model, prompt_tokens, completion_tokens):
"""Record token usage""" """Record token usage"""
today = date.today() today = date.today()