Spring AI 代码分析（五）–RAG 分析

技术分享 4个月前 (11-23) 0 999+

Spring AI RAG 分析

请关注微信公众号：阿呆-bot

1. 工程结构概览

spring-ai-rag 是 Spring AI 的检索增强生成（Retrieval Augmented Generation）框架，它提供了完整的 RAG 能力，让 AI 模型能够访问外部知识库。

spring-ai-rag/ ├── advisor/ │   └── RetrievalAugmentationAdvisor.java  # RAG Advisor │ ├── retrieval/                              # 检索层 │   └── search/ │       ├── DocumentRetriever.java          # 文档检索接口 │       └── VectorStoreDocumentRetriever.java  # 向量存储检索实现 │ ├── preretrieval/                           # 检索前处理 │   └── query/ │       ├── transformation/                 # 查询转换 │       │   ├── QueryTransformer.java │       │   ├── RewriteQueryTransformer.java │       │   └── CompressionQueryTransformer.java │       └── expansion/                       # 查询扩展 │           └── QueryExpander.java │ ├── postretrieval/                          # 检索后处理 │   └── DocumentPostProcessor.java │ └── generation/                             # 生成层     └── augmentation/         └── ContextualQueryAugmenter.java    # 上下文查询增强

2. 技术体系与模块关系

RAG 框架通过 Advisor 机制集成到 ChatClient，实现了完整的检索增强生成流程：

3. 关键场景示例代码

3.1 基础 RAG 使用

最简单的 RAG 使用方式：

@Autowired private ChatModel chatModel; @Autowired private VectorStore vectorStore;  public void basicRAG() {     // 创建向量存储检索器     VectorStoreDocumentRetriever retriever =          VectorStoreDocumentRetriever.builder()             .vectorStore(vectorStore)             .topK(5)             .similarityThreshold(0.7)             .build();          // 创建 RAG Advisor     RetrievalAugmentationAdvisor ragAdvisor =          RetrievalAugmentationAdvisor.builder()             .documentRetriever(retriever)             .build();          // 使用 ChatClient     ChatClient chatClient = ChatClient.builder(chatModel)         .defaultAdvisors(ragAdvisor)         .build();          String response = chatClient.prompt()         .user("查询文档中的信息")         .call()         .content(); }

3.2 带查询重写的 RAG

查询重写可以优化检索效果：

public void ragWithQueryRewrite() {     // 创建查询重写器     RewriteQueryTransformer queryTransformer =          RewriteQueryTransformer.builder()             .chatClient(ChatClient.builder(chatModel))             .targetSearchSystem("vector store")             .build();          // 创建 RAG Advisor     RetrievalAugmentationAdvisor ragAdvisor =          RetrievalAugmentationAdvisor.builder()             .documentRetriever(retriever)             .queryTransformers(queryTransformer)             .build();          ChatClient chatClient = ChatClient.builder(chatModel)         .defaultAdvisors(ragAdvisor)         .build(); }

3.3 多查询扩展

查询扩展可以从一个查询生成多个相关查询：

public void ragWithQueryExpansion() {     // 创建查询扩展器     QueryExpander queryExpander = new MultiQueryExpander(         ChatClient.builder(chatModel)     );          RetrievalAugmentationAdvisor ragAdvisor =          RetrievalAugmentationAdvisor.builder()             .documentRetriever(retriever)             .queryExpander(queryExpander)             .build(); }

4. 核心时序图

4.1 RAG 完整流程

5. 入口类与关键类关系

6. 关键实现逻辑分析

6.1 RetrievalAugmentationAdvisor 核心逻辑

RetrievalAugmentationAdvisor 是 RAG 流程的核心协调者：

public class RetrievalAugmentationAdvisor implements BaseAdvisor {     @Override     public ChatClientRequest before(ChatClientRequest request, AdvisorChain chain) {         // 1. 构建原始查询         Query originalQuery = Query.builder()             .text(request.prompt().getUserMessage().getText())             .history(request.prompt().getInstructions())             .context(request.context())             .build();                  // 2. 查询转换         Query transformedQuery = originalQuery;         for (QueryTransformer transformer : queryTransformers) {             transformedQuery = transformer.transform(transformedQuery);         }                  // 3. 查询扩展         List<Query> expandedQueries = queryExpander != null              ? queryExpander.expand(transformedQuery)             : List.of(transformedQuery);                  // 4. 并行检索文档         Map<Query, List<List<Document>>> documentsForQuery =              expandedQueries.parallelStream()                 .map(query -> getDocumentsForQuery(query))                 .collect(...);                  // 5. 合并文档         List<Document> documents = documentJoiner.join(documentsForQuery);                  // 6. 文档后处理         for (DocumentPostProcessor processor : documentPostProcessors) {             documents = processor.process(originalQuery, documents);         }                  // 7. 查询增强         Query augmentedQuery = queryAugmenter.augment(originalQuery, documents);                  // 8. 更新 Prompt         return request.mutate()             .prompt(augmentedQuery.toPrompt())             .context(context)             .build();     } }

6.2 向量存储检索实现

VectorStoreDocumentRetriever 负责从向量存储中检索文档：

public class VectorStoreDocumentRetriever implements DocumentRetriever {     @Override     public List<Document> retrieve(Query query) {         // 1. 计算请求的过滤表达式         Filter.Expression filterExpression = computeRequestFilterExpression(query);                  // 2. 构建搜索请求         SearchRequest searchRequest = SearchRequest.builder()             .query(query.text())             .filterExpression(filterExpression)             .similarityThreshold(this.similarityThreshold)             .topK(this.topK)             .build();                  // 3. 执行相似度搜索         return vectorStore.similaritySearch(searchRequest);     } }

6.3 查询重写实现

RewriteQueryTransformer 使用 LLM 重写查询以优化检索：

public class RewriteQueryTransformer implements QueryTransformer {     @Override     public Query transform(Query query) {         // 使用 LLM 重写查询         String rewrittenQueryText = chatClient.prompt()             .user(promptTemplate.getTemplate())             .param("target", targetSearchSystem)             .param("query", query.text())             .call()             .content();                  return query.mutate()             .text(rewrittenQueryText)             .build();     } }

6.4 查询增强实现

ContextualQueryAugmenter 将检索到的文档注入到用户查询中：

public class ContextualQueryAugmenter implements QueryAugmenter {     @Override     public Query augment(Query query, List<Document> documents) {         // 1. 格式化文档内容         String documentContext = formatDocuments(documents);                  // 2. 使用模板增强查询         String augmentedText = promptTemplate.render(Map.of(             "query", query.text(),             "context", documentContext         ));                  return query.mutate()             .text(augmentedText)             .build();     } }

7. 如何集成 Vector Store

RAG 框架通过 DocumentRetriever 接口抽象了向量存储的集成：

7.1 抽象机制

DocumentRetriever 接口提供了统一的文档检索抽象：

public interface DocumentRetriever extends Function<Query, List<Document>> {     List<Document> retrieve(Query query); }

这种设计让 RAG 框架不直接依赖 VectorStore，而是通过 DocumentRetriever 抽象。这意味着：

可以支持多种数据源：向量存储、知识图谱、数据库等
易于扩展：可以实现自定义的 DocumentRetriever
解耦：RAG 框架和向量存储完全解耦

7.2 VectorStoreDocumentRetriever 实现

VectorStoreDocumentRetriever 是向量存储的默认实现：

public class VectorStoreDocumentRetriever implements DocumentRetriever {     private final VectorStore vectorStore;          @Override     public List<Document> retrieve(Query query) {         SearchRequest request = SearchRequest.builder()             .query(query.text())             .topK(this.topK)             .similarityThreshold(this.similarityThreshold)             .filterExpression(computeFilterExpression(query))             .build();                  return vectorStore.similaritySearch(request);     } }

7.3 自定义 DocumentRetriever

可以轻松实现自定义的 DocumentRetriever：

public class CustomDocumentRetriever implements DocumentRetriever {     @Override     public List<Document> retrieve(Query query) {         // 自定义检索逻辑         // 可以从数据库、知识图谱、API 等检索         return customRetrieve(query);     } }