当前位置：首页 > news >正文

spring ai基于内存RAG尝鲜

news 2025/9/22 11:48:22

RAG，数据检索增强生成，简单点说你提供一个数据集，让语言模型根据你的数据集回答问题。
1.新增依赖
这次的练习demo是将一个pdf作为数据集，喂给模型做训练生成内存向量库，以此回答问题。下面新增的依赖是pdf阅读和内存向量库，spring-ai-transformers-spring-boot-starter中没有SimpleVectorStore，所以使用spring-ai-vector-store。

  	 <dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-pdf-document-reader</artifactId><version>1.0.0-SNAPSHOT</version></dependency><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-vector-store</artifactId><version>1.0.0-SNAPSHOT</version></dependency>

2.EmbeddingModel配置
EmbeddingModel用来完成"文本->向量集"的工作，包括数据集转成向量和把用户的问题转成向量。

import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.zhipuai.ZhiPuAiEmbeddingModel;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;@Configuration
public class VectorStoreConfiguration {@Beanpublic SimpleVectorStore vectorStore(ZhiPuAiEmbeddingModel zhiPuAiEmbeddingModel){return SimpleVectorStore.builder(zhiPuAiEmbeddingModel).build();}
}

3.加载向量集
项目启动时就需要加载向量集，第一次需要生成.json文件，后面就可以直接load加载，加快启动时间。

    @Autowiredprivate SimpleVectorStore simpleVectorStore; // 内存向量集@Autowiredprivate ResourceLoader resourceLoader;@PostConstruct@Overridepublic void init() {File file = new File("./vector_store.json");if (!file.exists()) {Resource pdfResource = resourceLoader.getResource("file:F:\\2024省政府工作报告.pdf");PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(pdfResource,PdfDocumentReaderConfig.builder().withPageExtractedTextFormatter(ExtractedTextFormatter.builder().build()).build());List<Document> documents = pdfReader.read();// 2. 分割文档TokenTextSplitter textSplitter = new TokenTextSplitter();List<Document> splitDocuments = textSplitter.split(documents);// 3. 嵌入并存储到向量集simpleVectorStore.add(splitDocuments);// 4. 保存到文件，下次启动无需重新处理simpleVectorStore.save(file);}// 加载simpleVectorStore.load(file);}

4.调用
4.1 根据向量集搜索是阻塞操作，但是我用的是webflux，导致报错，所以需要封装成响应式操作。
4.2 目前仅是练手demo，所以历史消息未做持久化存储和用户、话题隔离。

import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.zhipuai.ZhiPuAiChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;@RestController
public class DataController {@Autowiredprivate ZhiPuAiChatModel model;@Autowiredprivate SimpleVectorStore vectorStore;// todo 后续这里应改成Map<话题id,List<Message>>，这样子才能实现对话隔离List<Message> historyMessage = new ArrayList<>() {{add(new SystemMessage("你是一个乐于助人的AI助手，名字叫小智。你的回答要简洁明了。"));}};@GetMapping("/test/{msg}")public Flux<String> generate(@PathVariable(value = "msg") String msg) {// 1. 将阻塞的向量搜索转换为响应式操作Mono<List<Document>> similarDocumentsMono = Mono.fromCallable(() ->vectorStore.similaritySearch(msg)).subscribeOn(Schedulers.boundedElastic()); // 在弹性线程池执行阻塞操作// 2. 构建完整的响应式处理链return similarDocumentsMono.flatMapMany(documents -> {// 3. 构建系统消息和上下文String systemMessage = """请严格根据以下上下文信息来回答问题。如果上下文信息中没有答案，就说"根据我所知的信息，无法回答这个问题。"。不要编造信息。上下文信息：{context}""";String context = documents.stream().map(Document::getText).collect(Collectors.joining("\n\n"));Message systemDocument = new SystemPromptTemplate(systemMessage).createMessage(Map.of("context", context));// 4. 添加到历史消息（注意：这里需要处理历史消息的线程安全问题）// 异步环境中，需要确保 historyMessage 是线程安全的，后面改成用户和历史消息都隔离就不用加锁了synchronized (historyMessage) {historyMessage.add(systemDocument);Message userMessage = new UserMessage(msg);historyMessage.add(userMessage);}Prompt p = new Prompt(historyMessage);// 5. 调用流式模型Flux<String> responseStream = model.stream(p).handle((chatResponse, sink) -> {if (chatResponse.getResult() != null) {String text = chatResponse.getResult().getOutput().getText();if (text != null) {sink.next(text);}}});// 6. 处理流式响应并保存到历史return responseStream.collectList().flatMapMany(chunks -> {String fullResponse = String.join("", chunks);Message assistantMessage = new AssistantMessage(fullResponse);// 同样需要同步处理共享的historyMessagesynchronized (historyMessage) {historyMessage.add(assistantMessage);}return Flux.fromIterable(chunks);});});}
}

5.结果
GET http://localhost:8080/test/2024省政府工作报告可以提供几道题么？

回答：
根据上下文信息，我可以为您提供以下几道关于2024年山东省政府工作报告的题目：

（单选）《2024年山东省政府工作报告》围绕塑造绿色低碳高质量发展新优势，提出2024年要重点抓实抓好十二个方面的工作，其中放在首位的是（）。
A. 抓实抓好扩大有效需求
B. 科技创新引领现代化产业体系建设
C. 推进乡村振兴
D. 深化改革开放
（单选）根据2024年山东省政府工作报告，今年是中华人民共和国成立（）周年，是实现"十四五"规划目标任务的关键一年。
A. 70
B. 75
C. 80
D. 85
（多选）2024年山东省经济社会发展主要预期目标包括（）。
A. 地区生产总值增长5%以上
B. 一般公共预算收入增长4%
C. 居民人均可支配收入增长5.5%左右
D. 城镇新增就业100万人以上
（填空）2024年山东省政府工作以建设____________________为总抓手。