基于JWT的多租户RAG技术实现解析

近年来，大型语言模型（LLM）的出现加速了AI在各行业的应用。然而，为了进一步增强LLM能力并有效利用最新信息和领域知识，与外部数据源的集成至关重要。检索增强生成（RAG）技术因此受到广泛关注。

RAG通过从现有知识库检索相关信息，将其整合到LLM输入中生成更准确的响应。该技术已应用于产品开发的技术文档查询、客户支持FAQ应答，以及基于实时数据的决策系统。

在软件即服务（SaaS）场景中，多租户架构要求从单一代码库为多个租户提供服务。当使用LLM实现AI功能时，RAG技术能利用各租户特定数据提供个性化服务。以客服呼叫中心SaaS为例，每个租户的历史咨询记录、FAQ和产品手册构成其专属知识库，RAG系统可据此生成符合租户业务场景的精准响应。

然而，从安全角度考虑，多租户环境下的数据隔离成为关键挑战。某中心的Bedrock知识库简化了RAG实现，当选择OpenSearch作为向量数据库时，存在两种方案：

• OpenSearch无服务器版：支持通过元数据过滤实现租户隔离，但写入权限未分离
• OpenSearch服务版：具备细粒度访问控制（FGAC），但知识库仅支持单一IAM角色访问

本文提出的解决方案结合JSON Web Token（JWT）与FGAC实现租户隔离，主要技术优势包括：

1. 动态租户识别：JWT负载包含租户上下文属性，系统可动态识别请求所属租户
2. OpenSearch FGAC集成：直接利用JWT中的属性信息进行角色映射，实现索引或文档级访问控制

解决方案架构

系统采用OpenSearch作为向量数据库，整体流程如下：

1. 租户用户在Amazon Cognito用户池中创建，登录时通过Lambda触发器将租户ID注入JWT
2. 用户查询通过API Gateway传递至Lambda，附带JWT令牌
3. 查询文本通过某中心的文本嵌入模型向量化
4. 从DynamoDB获取目标OpenSearch域和索引信息
5. 在OpenSearch执行向量搜索，将结果作为上下文输入LLM生成响应

关键实现包含三个数据隔离模式：

• 域级隔离：每个租户使用独立OpenSearch域，FGAC角色授予全索引访问权限
• 索引级隔离：多租户共享OpenSearch域，FGAC角色限制仅能访问特定租户索引
• 文档级隔离：多租户共享域和索引，通过FGAC文档级安全查询实现数据隔离

技术实现细节

1. JWT生成与验证：通过Cognito用户池预生成令牌Lambda触发器，将DynamoDB存储的租户信息注入JWT
2. 请求路由：解析JWT获取租户ID，查询DynamoDB路由表确定目标OpenSearch端点
3. FGAC配置：在OpenSearch中创建租户专属角色，将JWT中的租户ID属性映射为后端角色
4. 隔离模式实现：
   • 文档级隔离通过dls参数设置租户ID过滤条件
   • 索引级隔离限制index_patterns为租户特定索引
   • 域级隔离通过独立域名实现物理隔离

生产环境建议

• 考虑采用DynamoDB分区模型优化多租户数据存储
• 可结合动态IAM策略增强资源访问控制
• 根据隔离严格性、管理复杂度等需求选择合适的隔离模式

该方案通过AWS CDK实现基础设施即代码部署，完整代码已开源在GitHub仓库。清理资源时只需执行cdk destroy命令即可删除所有关联资源。

本文展示的方案适用于需要严格租户数据隔离的SaaS场景，开发者可根据实际需求选择三种隔离模式中的最优方案。更多多租户RAG实现方案可参考某机构技术文档中的相关资源。