手搓生产级 AI Agent 系统（2）：工具系统深度实战——从参数校验到沙箱安全

import inspect
import json
import time
import signal
import re
from typing import Callable, Any, Optional, get_type_hints
from dataclasses import dataclass, field
from functools import wraps

# ====== 全局工具注册表 ======
TOOL_REGISTRY: dict[str, 'ToolDefinition'] = {}


@dataclass
class ToolParameter:
    """工具参数定义——给 LLM 看的描述直接决定调用准确率"""
    name: str
    type: str                    # string / number / boolean / object
    description: str             # 这个描述是给 LLM 看的，写清楚传什么
    required: bool = True
    default: Any = None
    enum: Optional[list] = None  # 限定可选值，大幅降低参数错误率


@dataclass
class ToolDefinition:
    """工具的完整定义"""
    name: str
    description: str             # 包含"做什么"和"不要做什么"
    parameters: list[ToolParameter]
    func: Callable
    category: str = "general"    # 工具分类，方便组织
    require_confirmation: bool = False  # 危险操作需用户确认
    timeout: int = 30            # 超时秒数
    max_retries: int = 1         # 失败重试次数
    risk_level: str = "safe"     # safe / dangerous / destructive


def register_tool(
    name: str,
    description: str,
    parameters: list[dict] = None,
    category: str = "general",
    require_confirmation: bool = False,
    timeout: int = 30,
    max_retries: int = 1,
    risk_level: str = "safe",
):
    """装饰器：将任意 Python 函数注册为 Agent 工具。

    用法：
        @register_tool(
            name="send_email",
            description="发送邮件。参数 to 为收件人邮箱...",
            parameters=[
                {"name": "to", "type": "string", "description": "收件人邮箱地址", "required": True},
                {"name": "subject", "type": "string", "description": "邮件主题", "required": True},
                {"name": "body", "type": "string", "description": "邮件正文（支持 Markdown）", "required": True},
            ],
            category="communication",
            require_confirmation=True,
            risk_level="dangerous",
        )
        def send_email(to: str, subject: str, body: str) -> str:
            ...
    """
    if parameters is None:
        parameters = []

    def decorator(func: Callable) -> Callable:
        # 从函数签名自动推断参数类型（如果没在 parameters 里显式声明）
        hints = get_type_hints(func) if hasattr(func, '__annotations__') else {}
        param_list = []
        for p in parameters:
            param_type = p.get("type", "string")
            # 如果类型没指定，从 type hints 推断
            if "type" not in p or p["type"] == "string":
                py_type = hints.get(p["name"])
                if py_type == int:
                    param_type = "number"
                elif py_type == bool:
                    param_type = "boolean"
                elif py_type in (list, dict):
                    param_type = "object"
            param_list.append(ToolParameter(
                name=p["name"],
                type=param_type,
                description=p.get("description", ""),
                required=p.get("required", True),
                default=p.get("default"),
                enum=p.get("enum"),
            ))

        tool_def = ToolDefinition(
            name=name,
            description=description,
            parameters=param_list,
            func=func,
            category=category,
            require_confirmation=require_confirmation,
            timeout=timeout,
            max_retries=max_retries,
            risk_level=risk_level,
        )
        TOOL_REGISTRY[name] = tool_def
        return func

    return decorator


def get_tools_for_llm(category: str = None) -> str:
    """生成给 LLM 看的工具描述。格式直接影响调用准确率。"""
    tools = TOOL_REGISTRY.values()
    if category:
        tools = [t for t in tools if t.category == category]

    lines = []
    for tool in tools:
        params_desc = []
        for p in tool.parameters:
            desc = f"    - {p.name} ({p.type}): {p.description}"
            if p.required:
                desc += " [必填]"
            if p.enum:
                desc += f" 可选值: {p.enum}"
            params_desc.append(desc)

        lines.append(f"""### {tool.name}
{tool.description}

参数：
{chr(10).join(params_desc)}
风险等级：{tool.risk_level}
""")

    return "\n".join(lines)

class ParameterValidator:
    """参数校验器——在工具执行前校验所有参数"""

    @staticmethod
    def validate(tool: ToolDefinition, params: dict) -> tuple[bool, str]:
        """校验参数。返回 (是否通过, 错误信息)。"""
        # 1. 必填参数检查
        for p in tool.parameters:
            if p.required and p.name not in params:
                return False, f"缺少必填参数 '{p.name}'。请提供 {p.description}"

        # 2. 未知参数检查
        valid_names = {p.name for p in tool.parameters}
        for key in params:
            if key not in valid_names:
                return False, f"未知参数 '{key}'。{tool.name} 接受的参数: {list(valid_names)}"

        # 3. 类型校验 + 自动转换
        for p in tool.parameters:
            if p.name not in params:
                continue
            value = params[p.name]

            if p.type == "number":
                if not isinstance(value, (int, float)):
                    try:
                        params[p.name] = float(value)  # 自动转换字符串数字
                    except (ValueError, TypeError):
                        return False, f"参数 '{p.name}' 需要数字，收到了 '{value}'"

            elif p.type == "boolean":
                if isinstance(value, str):
                    if value.lower() in ("true", "yes", "1"):
                        params[p.name] = True
                    elif value.lower() in ("false", "no", "0", ""):
                        params[p.name] = False
                    else:
                        return False, f"参数 '{p.name}' 需要布尔值(true/false)，收到了 '{value}'"

            elif p.type == "object":
                if isinstance(value, str):
                    try:
                        params[p.name] = json.loads(value)
                    except json.JSONDecodeError:
                        return False, f"参数 '{p.name}' 需要 JSON 对象，收到了 '{value[:50]}...'"

            # 4. 枚举值校验
            if p.enum and params[p.name] not in p.enum:
                return False, f"参数 '{p.name}' 的值 '{params[p.name]}' 不在允许范围 {p.enum}"

        return True, "OK"

┌─────────────────────────────────────────┐
│           安全沙箱三层防护                │
├─────────────────────────────────────────┤
│ Layer 1: 输入层                          │
│ - 参数校验（类型/必填/枚举）              │
│ - 危险模式检测（rm -rf, DROP TABLE...）  │
│ - SQL 注入 / Shell 注入检测              │
├─────────────────────────────────────────┤
│ Layer 2: 执行层                          │
│ - risk_level == "destructive" → 需要确认 │
│ - 超时控制（signal.alarm）               │
│ - 文件操作限制在安全目录                  │
├─────────────────────────────────────────┤
│ Layer 3: 审计层                          │
│ - 每次调用记录完整日志                    │
│ - 破坏性操作的参数持久化保存              │
│ - 异常告警（webhook/邮件）               │
└─────────────────────────────────────────┘

class SecureToolExecutor:
    """安全的工具执行器"""

    # 危险操作模式——正则匹配拦截
    DANGEROUS_PATTERNS = [
        (r'rm\s+(-rf?|--recursive)', '危险的文件删除操作'),
        (r'>\s*/dev/', '危险的设备写入'),
        (r'mkfs\.', '格式化文件系统'),
        (r'dd\s+if=', '磁盘直接读写'),
        (r'chmod\s+777', '危险的权限修改'),
        (r'(?i)DROP\s+TABLE', '数据库删表操作'),
        (r'(?i)DELETE\s+FROM\s+\w+\s*$', '数据库无条件删除'),
        (r'(?i)UPDATE\s+\w+\s+SET.*WHERE\s*$', '数据库无条件更新'),
    ]

    @staticmethod
    def check_dangerous(params: dict) -> tuple[bool, str]:
        """检查参数中是否包含危险操作。"""
        params_str = json.dumps(params, ensure_ascii=False)
        for pattern, description in SecureToolExecutor.DANGEROUS_PATTERNS:
            if re.search(pattern, params_str):
                return True, f"检测到危险操作: {description}"
        return False, ""

    @staticmethod
    def execute(tool: ToolDefinition, params: dict) -> dict:
        """执行工具调用——完整的防护链路。"""
        tracer = ToolTracer(tool.name)

        # Layer 1: 参数校验
        valid, error = ParameterValidator.validate(tool, params)
        if not valid:
            tracer.log_error("validation_error", error)
            return {
                "success": False,
                "error": f"参数校验失败: {error}",
                "suggestion": "请修正参数后重试",
            }

        # Layer 1: 危险模式检测
        is_dangerous, danger_desc = SecureToolExecutor.check_dangerous(params)
        if is_dangerous and tool.risk_level != "destructive":
            tracer.log_error("dangerous_blocked", danger_desc)
            return {
                "success": False,
                "error": f"操作被拦截: {danger_desc}",
                "blocked": True,
            }

        # Layer 2: 需要确认的危险操作
        if tool.require_confirmation:
            return {
                "success": False,
                "error": f"操作 '{tool.name}' 需要用户确认",
                "require_confirmation": True,
                "risk_level": tool.risk_level,
            }

        # Layer 2: 执行（带超时 + 重试）
        for attempt in range(tool.max_retries + 1):
            try:
                tracer.start()
                result = tool.func(**params)
                tracer.end(len(str(result)))
                tracer.log_success(result)
                return {"success": True, "result": result, "attempt": attempt + 1}

            except Exception as e:
                if attempt == tool.max_retries:
                    tracer.log_error(type(e).__name__, str(e))
                    return {
                        "success": False,
                        "error": f"执行失败（{attempt+1}/{tool.max_retries+1}次）: {str(e)}",
                        "attempts": attempt + 1,
                        "suggestion": "请尝试其他方式或告知用户此功能暂时不可用",
                    }
                time.sleep(0.5 * (attempt + 1))  # 指数退避

class ToolTracer:
    """每次工具调用的完整追踪记录"""

    _traces: list[dict] = []

    def __init__(self, tool_name: str):
        self.tool_name = tool_name
        self.start_time = 0
        self.end_time = 0

    def start(self):
        self.start_time = time.time()

    def end(self, result_size: int = 0):
        self.end_time = time.time()

    def log_success(self, result):
        ToolTracer._traces.append({
            "tool": self.tool_name,
            "status": "success",
            "latency_ms": int((self.end_time - self.start_time) * 1000),
            "result_preview": str(result)[:200],
            "timestamp": int(time.time()),
        })

    def log_error(self, error_type: str, detail: str):
        ToolTracer._traces.append({
            "tool": self.tool_name,
            "status": "error",
            "error_type": error_type,
            "error_detail": detail,
            "timestamp": int(time.time()),
        })

    @classmethod
    def get_summary(cls) -> dict:
        traces = cls._traces
        if not traces:
            return {}
        success = [t for t in traces if t["status"] == "success"]
        errors = [t for t in traces if t["status"] == "error"]
        latencies = [t.get("latency_ms", 0) for t in success]
        return {
            "total_calls": len(traces),
            "success_rate": len(success) / len(traces) * 100,
            "avg_latency_ms": sum(latencies) / len(latencies) if latencies else 0,
            "error_breakdown": {
                et: len([t for t in errors if t.get("error_type") == et])
                for et in set(t.get("error_type") for t in errors)
            }
        }

# ====== 注册工具 ======

@register_tool(
    name="web_search",
    description="搜索互联网获取实时信息。用于查询最新新闻、事实数据、人物信息。不要用于数学计算、代码审查、文件操作。",
    parameters=[
        {"name": "query", "type": "string",
         "description": "搜索关键词，应包含核心术语。例如查'马斯克年龄'而非'那个特斯拉老板多大了'",
         "required": True},
        {"name": "max_results", "type": "number",
         "description": "返回结果数量，默认5，最大10",
         "required": False, "default": 5},
    ],
    category="search"
)
def web_search(query: str, max_results: int = 5) -> str:
    # 生产环境替换为真实搜索 API
    return f'搜索"{query}"结果（{max_results}条）: ...'


@register_tool(
    name="calculator",
    description="执行数学计算。支持加减乘除、括号、幂运算。用于需要精确计算的场景。",
    parameters=[
        {"name": "expression", "type": "string",
         "description": "数学表达式，如 '(2+3)*4' 或 'sqrt(16)'",
         "required": True},
    ],
    category="utility"
)
def calculator(expression: str) -> str:
    # 安全沙箱：只允许安全的数学运算
    allowed = set("0123456789+-*/.() eExXpiPIsqrtanlgo ")
    if not all(c in allowed for c in expression):
        return f"错误：表达式包含不允许的字符"
    try:
        result = eval(expression, {"__builtins__": {}}, {
            "sqrt": lambda x: x ** 0.5, "pi": 3.1415926535, "e": 2.718281828,
            "abs": abs, "round": round, "min": min, "max": max,
        })
        return f"计算结果: {result}"
    except Exception as e:
        return f"计算失败: {e}"


@register_tool(
    name="read_file",
    description="读取指定文件内容。用于查看代码、配置、日志等文本文件。不要用来读取二进制文件（图片、视频、可执行文件）。",
    parameters=[
        {"name": "path", "type": "string",
         "description": "文件路径，相对于项目根目录",
         "required": True},
    ],
    category="file",
    risk_level="dangerous",
)
def read_file(path: str) -> str:
    # 路径校验：防止目录遍历攻击
    if ".." in path or path.startswith("/"):
        return "错误：不允许访问上级目录或绝对路径"
    if not os.path.exists(path):
        return f"错误：文件 {path} 不存在"
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()[:3000]  # 限制读取长度
    return content


@register_tool(
    name="delete_file",
    description="删除指定文件。此操作不可逆，请谨慎使用。删除前建议先确认文件内容。",
    parameters=[
        {"name": "path", "type": "string",
         "description": "要删除的文件路径",
         "required": True},
    ],
    category="file",
    require_confirmation=True,
    risk_level="destructive",
)
def delete_file(path: str) -> str:
    if not os.path.exists(path):
        return f"文件 {path} 不存在"
    os.remove(path)
    return f"已删除: {path}"


@register_tool(
    name="get_current_time",
    description="获取当前日期和时间。当用户问'今天''现在''当前时间'时使用。",
    parameters=[],
    category="utility"
)
def get_current_time() -> str:
    return time.strftime("当前时间：%Y年%m月%d日 %H:%M:%S", time.localtime())


# ====== 测试工具调用准确率 ======
if __name__ == "__main__":
    print(f"注册工具: {len(TOOL_REGISTRY)} 个")
    print(get_tools_for_llm())

    # 测试 safe 工具
    executor = SecureToolExecutor()
    result = executor.execute(
        TOOL_REGISTRY["calculator"],
        {"expression": "2 + 3 * 4"}
    )
    print(f"calculator: {result}")

    # 测试 destructive 工具（需要确认）
    result = executor.execute(
        TOOL_REGISTRY["delete_file"],
        {"path": "test.txt"}
    )
    print(f"delete_file: {result}")