引言:为什么Agent的工具调用需要“双重保险”?
在AI Agent的生产部署中,工具调用(Tool Use)是核心功能,但也是最大的安全隐患和可靠性挑战。LLM可能会因幻觉(Hallucination)生成无效、格式错误的参数,或者尝试调用用户无权访问的敏感功能(权限逃逸)。
为了解决这些问题,我们需要一个在模型输出和实际执行之间工作的安全层。本文将介绍如何结合使用Python生态中最强大的数据验证库 Pydantic 进行结构化输入验证,并通过自定义装饰器实现灵活的运行时权限(RBAC/ABAC)控制。
第一步:使用Pydantic定义工具的严格输入Schema
Agent模型通常会输出一个JSON字符串作为工具调用的参数。在执行函数之前,我们必须确保这个JSON不仅格式正确,而且字段类型、必填性都符合预期。
Pydantic是实现这一目标的最佳选择,它允许我们以声明式的方式定义工具的参数结构。
1.1 定义Pydantic参数模型
假设我们有一个敏感的工具 file_writer,它接受 path 和 content 两个参数。
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 系统的稳定性和安全性。
汤不热吧