欢迎光临
我们一直在努力

如何为Agent的工具调用(Tool Use)实现严格的输入验证和权限控制?

引言:为什么Agent的工具调用需要“双重保险”?

在AI Agent的生产部署中,工具调用(Tool Use)是核心功能,但也是最大的安全隐患和可靠性挑战。LLM可能会因幻觉(Hallucination)生成无效、格式错误的参数,或者尝试调用用户无权访问的敏感功能(权限逃逸)。

为了解决这些问题,我们需要一个在模型输出和实际执行之间工作的安全层。本文将介绍如何结合使用Python生态中最强大的数据验证库 Pydantic 进行结构化输入验证,并通过自定义装饰器实现灵活的运行时权限(RBAC/ABAC)控制。

第一步:使用Pydantic定义工具的严格输入Schema

Agent模型通常会输出一个JSON字符串作为工具调用的参数。在执行函数之前,我们必须确保这个JSON不仅格式正确,而且字段类型、必填性都符合预期。

Pydantic是实现这一目标的最佳选择,它允许我们以声明式的方式定义工具的参数结构。

1.1 定义Pydantic参数模型

假设我们有一个敏感的工具 file_writer,它接受 pathcontent 两个参数。

from pydantic import BaseModel, Field, ValidationError
import json
from typing import Dict, Any

# 1. 定义工具参数的严格Schema
class FileSystemToolArgs(BaseModel):
    """用于写入文件的工具参数定义"""
    path: str = Field(..., description="要操作的文件路径,必须是绝对路径。")
    content: str = Field(..., description="写入文件的内容,不能为空。")

# 2. 通用的输入验证函数
def validate_tool_input(raw_args_json: str, schema: BaseModel) -> Dict[str, Any]:
    """解析LLM输出的JSON,并依据Pydantic Schema进行验证。"""
    try:
        # 尝试解析JSON
        args_dict = json.loads(raw_args_json)
        # 实例化Schema,触发验证
        validated_data = schema(**args_dict)
        return validated_data.dict()
    except json.JSONDecodeError:
        raise ValueError("输入不是有效的JSON格式,拒绝执行。")
    except ValidationError as e:
        # Pydantic 验证失败,返回给LLM进行修正或终止执行
        raise ValueError(f"参数验证失败,请检查字段类型或必填项: {e}")

# 示例:尝试验证一个错误的输入
# try:
#     invalid_input = '{"path": "/tmp/log.txt"}' # 缺少content字段
#     validate_tool_input(invalid_input, FileSystemToolArgs)
# except ValueError as e:
#     print(f"验证结果: {e}")

通过这一步,我们成功过滤掉了结构性错误和类型不匹配的参数,大大提高了工具的健壮性。

第二步:实现基于装饰器的权限控制(Access Control)

即使输入参数通过了验证,我们还需要确保执行工具的上下文(例如用户Session、Agent ID)拥有足够的权限。这里,我们使用Python装饰器(Decorator)实现运行时权限检查,这是一种干净且可重用的方式。

我们将权限控制逻辑与实际的工具函数解耦。

2.1 定义权限检查装饰器

import functools
from typing import Callable, Any

# 模拟权限系统:定义不同角色的可用工具集
SYSTEM_PERMISSIONS = {
    "admin": ["write_file", "read_db", "shutdown_server"],
    "editor": ["write_file", "read_db"],
    "guest": ["read_db"]
}

def requires_permission(tool_name: str):
    """权限检查装饰器,用于包裹工具函数"""
    def decorator(func: Callable):
        @functools.wraps(func)
        def wrapper(user_role: str, *args, **kwargs) -> Any:
            # 检查用户角色是否拥有调用该工具的权限
            if tool_name not in SYSTEM_PERMISSIONS.get(user_role, []):
                # 权限不足,抛出异常,阻止工具执行
                raise PermissionError(f"用户角色 '{user_role}' 无权使用工具 '{tool_name}'。")

            print(f"[ACCESS GRANTED] User '{user_role}' executing '{tool_name}'.")
            # 权限通过,执行原始函数
            return func(*args, **kwargs)
        return wrapper
    return decorator

# 2.2 应用权限控制到工具函数
@requires_permission(tool_name="write_file")
def safe_write_file(user_role: str, path: str, content: str):
    # 注意:user_role参数是必需的,用于装饰器内部的权限检查
    print(f"执行文件写入操作:Path={path}, Content Length={len(content)}")
    return f"Successfully written by {user_role} to {path}"

第三步:集成Agent工具调用流程

现在我们将 Pydantic 验证和权限控制结合起来,构建一个健壮的工具执行管道(Pipeline)。

TOOL_MAP = {
    "write_file": {
        "func": safe_write_file,
        "schema": FileSystemToolArgs
    }
}

def execute_agent_tool_call(
    user_role: str, 
    tool_name: str, 
    raw_args_json: str
):
    """Agent工具调用的安全执行入口"""

    if tool_name not in TOOL_MAP:
        return {"error": "Tool not found.", "code": 404}

    tool_config = TOOL_MAP[tool_name]

    print(f"\n--- Attempting Tool Call: {tool_name} by {user_role} ---")

    # Step 1: 严格参数验证 (Pydantic)
    try:
        validated_args = validate_tool_input(raw_args_json, tool_config["schema"])
        print("[Validation Success]")
    except ValueError as e:
        # LLM输出参数错误
        return {"error": f"Validation Failed. {str(e)}", "code": 400}

    # Step 2: 权限控制和执行 (Decorator enforced)
    try:
        # 传入 user_role 进行权限检查
        result = tool_config["func"](user_role=user_role, **validated_args)
        return {"status": "success", "result": result, "code": 200}
    except PermissionError as e:
        # 用户权限不足
        return {"error": f"Permission Denied. {str(e)}", "code": 403}

# --- 实际操作示例 ---

# 示例 A: Admin用户尝试合法调用 (应成功)
admin_input = '{"path": "/critical/config.yml", "content": "update"}'
result_a = execute_agent_tool_call("admin", "write_file", admin_input)
print(result_a)

# 示例 B: Guest用户尝试调用敏感工具 (应因权限失败)
guest_input = '{"path": "/tmp/log.txt", "content": "log data"}'
result_b = execute_agent_tool_call("guest", "write_file", guest_input)
print(result_b)

# 示例 C: Admin用户传入结构错误参数 (应因Pydantic验证失败)
err_input = '{"path": 12345, "content": "test"}' # path应为字符串
result_c = execute_agent_tool_call("admin", "write_file", err_input)
print(result_c)

运行结果预期:

  • 示例 A 成功执行。
  • 示例 B 在权限检查阶段被阻止,返回 Code 403。
  • 示例 C 在 Pydantic 验证阶段被阻止,返回 Code 400,并指出 path 字段类型错误。

总结

为 Agent 工具调用引入严格的输入验证和权限控制是构建生产级安全 Agent 基础设施的关键一步。通过将 Pydantic 用于精确的参数结构化,以及自定义 Python 装饰器 实现灵活的访问控制策略,我们能够有效隔离 LLM 产生的不可靠输出,防止未授权的操作,从而确保整个 Agent 系统的稳定性和安全性。

【本站文章皆为原创,未经允许不得转载】:汤不热吧 » 如何为Agent的工具调用(Tool Use)实现严格的输入验证和权限控制?
分享到: 更多 (0)

评论 抢沙发

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址