欢迎光临
我们一直在努力

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

引言

在构建基于大语言模型(LLM)的 Agent 时,工具调用(Tool Use)是实现模型与外部系统交互的核心。然而,LLM 输出的不确定性极易引发安全风险,例如模型可能会生成超出范围的参数,或者尝试调用超出其权限的操作。本文将介绍如何通过 Pydantic 强校验与动态权限中间件,为 Agent 工具库构建一套完善的防御体系。

1. 使用 Pydantic 构建强类型约束

第一道防线是定义严格的输入 Schema。仅靠模型理解文档是不够的,我们需要在执行层面拦截不合规的参数。

from pydantic import BaseModel, Field, validator
from typing import Optional, List

class FileOperationSchema(BaseModel):
    operation: str = Field(..., description=\"操作类型: read, write\")
    file_path: str = Field(..., description=\"相对于安全目录的文件路径\")
    content: Optional[str] = Field(None, description=\"写入的内容\")

    @validator('file_path')
    def prevent_path_traversal(cls, v):
        if '..' in v or v.startswith('/'):
            raise ValueError(\"非法路径:禁止路径遍历或绝对路径访问\")
        return v

    @validator('operation')
    def validate_op(cls, v):
        if v not in ['read', 'write']:
            raise ValueError(\"不支持的操作类型\")
        return v

2. 实现权限管理中间件

即使输入合法,模型也可能被诱导执行当前用户无权执行的操作。我们需要引入一个 SecurityContext 来进行运行时权限校验。

from functools import wraps

class SecurityContext:
    def __init__(self, user_id: str, permissions: List[str]):
        self.user_id = user_id
        self.permissions = permissions

def requires_permission(permission_name: str):
    def decorator(func):
        @wraps(func)
        def wrapper(context: SecurityContext, **kwargs):
            if permission_name not in context.permissions:
                return {\"error\": f\"用户 {context.user_id} 缺少权限: {permission_name}\"}
            return func(**kwargs)
        return wrapper
    return decorator

3. 工具调用的完整实现

我们将上述验证与权限控制整合到一个工具执行器中。

@requires_permission(\"filesystem_access\")
def safe_file_tool(operation: str, file_path: str, content: Optional[str] = None):
    # 这里执行实际的 IO 操作
    return {\"status\": \"success\", \"msg\": f\"{operation} on {file_path} completed.\"}

def call_agent_tool(user_ctx: SecurityContext, tool_input: dict):
    try:
        # 1. 输入验证 (Schema Validation)
        validated_data = FileOperationSchema(**tool_input)

        # 2. 权限验证与执行 (Permission & Execution)
        # 注意:这里假设通过某种映射找到了对应的函数
        result = safe_file_tool(user_ctx, **validated_data.dict())
        return result
    except Exception as e:
        return {\"status\": \"error\", \"message\": str(e)}

# 示例调用
ctx = SecurityContext(user_id=\"user_001\", permissions=[]) # 无权限
invalid_input = {\"operation\": \"write\", \"file_path\": \"../../etc/passwd\", \"content\": \"hack\"}

print(call_agent_tool(ctx, invalid_input))
# 输出: {'status': 'error', 'message': '1 validation error for FileOperationSchema\
file_path\
  非法路径:禁止路径遍历或绝对路径访问 (type=value_error)'}

4. 最佳实践建议

  1. 沙箱化环境:对于执行代码或复杂命令的工具,必须在 Docker 容器或 WebAssembly 沙箱中运行。
  2. 审计日志:记录每一次工具调用的输入、输出、用户上下文及权限校验结果,以便于溯源。
  3. 最小权限原则:不要给 Agent 赋予 root 或 admin 权限,应根据业务场景动态分配精细化的 Scope。

总结

通过将 Pydantic 的静态约束与基于装饰器的动态权限校验结合,我们可以有效地降低 LLM Agent 在调用外部工具时的攻击面。这不仅保护了系统安全,也提高了 Agent 在企业级生产环境中的可靠性。

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

评论 抢沙发

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