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

Spring AI RAG 分析

请关注微信公众号：阿呆-bot
1. 工程结构概览

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

1. spring-ai-rag/
2. ├── advisor/
3. │   └── RetrievalAugmentationAdvisor.java  # RAG Advisor
4. │
5. ├── retrieval/                              # 检索层
6. │   └── search/
7. │       ├── DocumentRetriever.java          # 文档检索接口
8. │       └── VectorStoreDocumentRetriever.java  # 向量存储检索实现
9. │
10. ├── preretrieval/                           # 检索前处理
11. │   └── query/
12. │       ├── transformation/                 # 查询转换
13. │       │   ├── QueryTransformer.java
14. │       │   ├── RewriteQueryTransformer.java
15. │       │   └── CompressionQueryTransformer.java
16. │       └── expansion/                       # 查询扩展
17. │           └── QueryExpander.java
18. │
19. ├── postretrieval/                          # 检索后处理
20. │   └── DocumentPostProcessor.java
21. │
22. └── generation/                             # 生成层
23.     └── augmentation/
24.         └── ContextualQueryAugmenter.java    # 上下文查询增强

复制代码

2. 技术体系与模块关系

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

3. 关键场景示例代码

3.1 基础 RAG 使用

最简单的 RAG 使用方式：

1. @Autowired
2. private ChatModel chatModel;
3. @Autowired
4. private VectorStore vectorStore;
6. public void basicRAG() {
7.     // 创建向量存储检索器
8.     VectorStoreDocumentRetriever retriever =
9.         VectorStoreDocumentRetriever.builder()
10.             .vectorStore(vectorStore)
11.             .topK(5)
12.             .similarityThreshold(0.7)
13.             .build();
14.    
15.     // 创建 RAG Advisor
16.     RetrievalAugmentationAdvisor ragAdvisor =
17.         RetrievalAugmentationAdvisor.builder()
18.             .documentRetriever(retriever)
19.             .build();
20.    
21.     // 使用 ChatClient
22.     ChatClient chatClient = ChatClient.builder(chatModel)
23.         .defaultAdvisors(ragAdvisor)
24.         .build();
25.    
26.     String response = chatClient.prompt()
27.         .user("查询文档中的信息")
28.         .call()
29.         .content();
30. }

复制代码

3.2 带查询重写的 RAG

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

1. public void ragWithQueryRewrite() {
2.     // 创建查询重写器
3.     RewriteQueryTransformer queryTransformer =
4.         RewriteQueryTransformer.builder()
5.             .chatClient(ChatClient.builder(chatModel))
6.             .targetSearchSystem("vector store")
7.             .build();
8.    
9.     // 创建 RAG Advisor
10.     RetrievalAugmentationAdvisor ragAdvisor =
11.         RetrievalAugmentationAdvisor.builder()
12.             .documentRetriever(retriever)
13.             .queryTransformers(queryTransformer)
14.             .build();
15.    
16.     ChatClient chatClient = ChatClient.builder(chatModel)
17.         .defaultAdvisors(ragAdvisor)
18.         .build();
19. }

复制代码

3.3 多查询扩展

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

1. public void ragWithQueryExpansion() {
2.     // 创建查询扩展器
3.     QueryExpander queryExpander = new MultiQueryExpander(
4.         ChatClient.builder(chatModel)
5.     );
6.    
7.     RetrievalAugmentationAdvisor ragAdvisor =
8.         RetrievalAugmentationAdvisor.builder()
9.             .documentRetriever(retriever)
10.             .queryExpander(queryExpander)
11.             .build();
12. }

复制代码

4. 核心时序图

4.1 RAG 完整流程

5. 入口类与关键类关系

6. 关键实现逻辑分析

6.1 RetrievalAugmentationAdvisor 核心逻辑

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

1. public class RetrievalAugmentationAdvisor implements BaseAdvisor {
2.     @Override
3.     public ChatClientRequest before(ChatClientRequest request, AdvisorChain chain) {
4.         // 1. 构建原始查询
5.         Query originalQuery = Query.builder()
6.             .text(request.prompt().getUserMessage().getText())
7.             .history(request.prompt().getInstructions())
8.             .context(request.context())
9.             .build();
10.         
11.         // 2. 查询转换
12.         Query transformedQuery = originalQuery;
13.         for (QueryTransformer transformer : queryTransformers) {
14.             transformedQuery = transformer.transform(transformedQuery);
15.         }
16.         
17.         // 3. 查询扩展
18.         List<Query> expandedQueries = queryExpander != null
19.             ? queryExpander.expand(transformedQuery)
20.             : List.of(transformedQuery);
21.         
22.         // 4. 并行检索文档
23.         Map<Query, List<List<Document>>> documentsForQuery =
24.             expandedQueries.parallelStream()
25.                 .map(query -> getDocumentsForQuery(query))
26.                 .collect(...);
27.         
28.         // 5. 合并文档
29.         List<Document> documents = documentJoiner.join(documentsForQuery);
30.         
31.         // 6. 文档后处理
32.         for (DocumentPostProcessor processor : documentPostProcessors) {
33.             documents = processor.process(originalQuery, documents);
34.         }
35.         
36.         // 7. 查询增强
37.         Query augmentedQuery = queryAugmenter.augment(originalQuery, documents);
38.         
39.         // 8. 更新 Prompt
40.         return request.mutate()
41.             .prompt(augmentedQuery.toPrompt())
42.             .context(context)
43.             .build();
44.     }
45. }

复制代码

6.2 向量存储检索实现

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

1. public class VectorStoreDocumentRetriever implements DocumentRetriever {
2.     @Override
3.     public List<Document> retrieve(Query query) {
4.         // 1. 计算请求的过滤表达式
5.         Filter.Expression filterExpression = computeRequestFilterExpression(query);
6.         
7.         // 2. 构建搜索请求
8.         SearchRequest searchRequest = SearchRequest.builder()
9.             .query(query.text())
10.             .filterExpression(filterExpression)
11.             .similarityThreshold(this.similarityThreshold)
12.             .topK(this.topK)
13.             .build();
14.         
15.         // 3. 执行相似度搜索
16.         return vectorStore.similaritySearch(searchRequest);
17.     }
18. }

复制代码

6.3 查询重写实现

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

1. public class RewriteQueryTransformer implements QueryTransformer {
2.     @Override
3.     public Query transform(Query query) {
4.         // 使用 LLM 重写查询
5.         String rewrittenQueryText = chatClient.prompt()
6.             .user(promptTemplate.getTemplate())
7.             .param("target", targetSearchSystem)
8.             .param("query", query.text())
9.             .call()
10.             .content();
11.         
12.         return query.mutate()
13.             .text(rewrittenQueryText)
14.             .build();
15.     }
16. }

复制代码

6.4 查询增强实现

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

1. public class ContextualQueryAugmenter implements QueryAugmenter {
2.     @Override
3.     public Query augment(Query query, List<Document> documents) {
4.         // 1. 格式化文档内容
5.         String documentContext = formatDocuments(documents);
6.         
7.         // 2. 使用模板增强查询
8.         String augmentedText = promptTemplate.render(Map.of(
9.             "query", query.text(),
10.             "context", documentContext
11.         ));
12.         
13.         return query.mutate()
14.             .text(augmentedText)
15.             .build();
16.     }
17. }

复制代码

7. 如何集成 Vector Store

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

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

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

复制代码

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

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

7.2 VectorStoreDocumentRetriever 实现

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

1. public class VectorStoreDocumentRetriever implements DocumentRetriever {
2.     private final VectorStore vectorStore;
3.    
4.     @Override
5.     public List<Document> retrieve(Query query) {
6.         SearchRequest request = SearchRequest.builder()
7.             .query(query.text())
8.             .topK(this.topK)
9.             .similarityThreshold(this.similarityThreshold)
10.             .filterExpression(computeFilterExpression(query))
11.             .build();
12.         
13.         return vectorStore.similaritySearch(request);
14.     }
15. }

复制代码

7.3 自定义 DocumentRetriever

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

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

复制代码

8. 外部依赖

spring-ai-rag 的依赖：

• spring-ai-client-chat：ChatClient 和 Advisor 机制
  
• spring-ai-vector-store：向量存储抽象
  
• spring-ai-commons：Document 等基础类型
  
• Spring Framework：IoC 和核心功能
  
• Reactor Core：响应式处理（用于并行检索）
  

9. 工程总结

Spring AI RAG 框架的设计有几个亮点：
模块化设计。RAG 流程被拆成了多个独立的组件：查询转换、查询扩展、文档检索、文档后处理、查询增强。每个组件都能独立配置和替换，想用哪个用哪个。
抽象和解耦。通过 DocumentRetriever 接口，RAG 框架和向量存储完全解耦。这样就能支持多种数据源，想换就换，扩展起来也方便。
Advisor 机制集成。RAG 通过 Advisor 机制集成到 ChatClient，用起来特别简单，只需要加一个 Advisor 就行。
可扩展性。每个环节都支持自定义实现，比如自己写个查询转换器、文档后处理器什么的。这样 RAG 框架就能适应各种业务场景。
并行处理。查询扩展后生成的多个查询可以并行检索，性能更好。
总的来说，Spring AI RAG 框架是一个设计得不错、功能完整的 RAG 实现。它提供了完整的 RAG 能力，同时保持了高度的灵活性和可扩展性。开发者可以轻松构建基于 RAG 的 AI 应用，也能根据具体需求进行深度定制。

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

很好很强大  我过来先占个楼 待编辑