基于Spring-AI框架实现RAG增强检索(附源码)

引言

随着人工智能技术的快速发展，增强检索（Retrieval-Augmented Generation, RAG）已成为一种结合检索和生成的先进方法，广泛应用于各种智能应用中。Spring-AI作为Spring Boot的AI扩展，提供了一套丰富的工具和库，使得开发者可以轻松地实现RAG技术。本文将介绍如何基于Spring-AI框架配置项目，并实现RAG增强检索，更多RAG增强检索的应用场景。

1. 实现效果

知识库文本-pet.txt：

客户姓名|宠物|洗澡日期|剪毛日期
张晓丽|加菲猫|2024年3月5日|2024年3月5日
张晓丽|泰迪犬|2024年6月18日|2024年5月18日
王宏|贵宾犬|2024年6月15日|2024年4月18日
王宏|阿拉斯加犬|2024年5月12日|2024年6月12日

知识库文本-pet-rule.txt

所有宠物超过15天需要洗澡一次，超过2个月需要剪毛。

项目启动时读取知识库文本文件，存入向量数据库，当调用接口时，先从向量数据库查出相关文本，然后将相关文本作为上下文连同问题，一并发给ai，最后ai依据上下文进行回答。

2. 准备工作

1. 安装ollama
2. 下载大模型

ollama pull wangshenzhi/llama3-8b-chinese-chat-ollama-q4
ollama pull nomic-embed-text

3. 项目配置（POM）

首先，我们需要在项目的pom.xml文件中引入Spring-AI相关的依赖。spring-boot-starter-webflux、spring-session-core等依赖，以及spring-ai-ollama和spring-ai-bom等Spring-AI核心库。

<dependencies><!-- Spring Boot WebFlux Starter --><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-webflux</artifactId></dependency><!-- Spring Session Core --><dependency><groupId>org.springframework.session</groupId><artifactId>spring-session-core</artifactId></dependency><!-- Spring AI Ollama Starter --><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-ollama</artifactId></dependency><!-- Spring AI BOM for Dependency Management --><dependencyManagement><dependencies><dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-bom</artifactId><version>1.0.0-SNAPSHOT</version><type>pom</type><scope>import</scope></dependency></dependencies></dependencyManagement></dependencies>

4. OllamaAutoConfig配置

接下来，我们需要配置OllamaAutoConfig类，该类负责初始化Ollama模型和相关组件。根据提供的文件，配置了OllamaApi、OllamaChatModel、OllamaEmbeddingModel和VectorStore。

@ConfigurationpublicclassOllamaAutoConfig{privatestaticfinalLogger log =LoggerFactory.getLogger(OllamaAutoConfig.class);@BeanpublicOllamaApiollamaApi(){returnnewOllamaApi("http://localhost:11434");}@BeanpublicOllamaChatModelchatModel(OllamaApi ollamaApi){OllamaChatModel chatModel =newOllamaChatModel(ollamaApi,OllamaOptions.create().withModel("wangshenzhi/llama3-8b-chinese-chat-ollama-q4").withTemperature(0.9f));return chatModel;}@BeanpublicOllamaEmbeddingModelembeddingModel(OllamaApi ollamaApi){returnnewOllamaEmbeddingModel(ollamaApi).withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));}@BeanpublicVectorStorevectorStore(OllamaEmbeddingModel embeddingModel){
        log.info("开始加载本地知识库文档");SimpleVectorStore vectorStore =newSimpleVectorStore(embeddingModel);File files =newFile(ClassLoader.getSystemResource("doc").getFile());for(File f:files.listFiles()){Resource resource =newFileSystemResource(f);TextReader textReader =newTextReader(resource);
            textReader.setCharset(Charset.defaultCharset());
            vectorStore.add(textReader.get());
            log.info("知识库文件加载完成：{}",f.getName());}return vectorStore;}}

5. Session配置

为了支持WebFlux的会话管理,同时让AI对话能记录聊天记录，这里将对话记录存到了session中，我们需要配置SessionConfig类。使用了ReactiveMapSessionRepository来存储会话信息。

@Configuration@EnableSpringWebSessionpublicclassSessionConfig{@BeanpublicReactiveSessionRepositoryreactiveSessionRepository(){returnnewReactiveMapSessionRepository(newConcurrentHashMap<>());}}

6. AI对话接口

这里使用了text/event-stream事件流式响应，从session中取出或初始化Prompt提示词对象，并将当前对话内容Message设置到Prompt中，这样对话过程就有了上下文，AI会依据上下文做出回答。

@GetMapping("/ai/generateStream")publicFlux<String>generateStream(@RequestParam(value ="message", defaultValue ="Tell me a joke")String message,ServerWebExchange exchange){ServerHttpResponse response = exchange.getResponse();
        response.getHeaders().add("Content-Type","text/event-stream");Mono<Prompt> promptMono = exchange.getSession().flatMap(webSession ->{Prompt prompt = webSession.getAttribute("prompt");if(prompt ==null){
                prompt =newPrompt(message);}else{List<Message> messages =newArrayList<>();
                messages.add(newUserMessage(prompt.getContents()));
                messages.add(newUserMessage(message));
                prompt =newPrompt(messages);}
            webSession.getAttributes().put("prompt",prompt);returnMono.just(prompt);});Flux<String> stream = promptMono.flatMapMany(this::getReply);return stream;}privateFlux<String>getReply(Prompt prompt){return chatModel.stream(prompt.getContents());}

7. 实现RAG增强检索

最后，我们来重点介绍如何在AIController中实现RAG增强检索。我们可以看到/ai/query接口的实现方法。

@GetMapping("/ai/query")publicFlux<String>query(@RequestParam(value ="message", defaultValue ="Tell me a joke")String message){// 使用自然语言查询 VectorStore，查找相关文档List<Document> similarDocuments = vectorStore.similaritySearch(SearchRequest.query(message).withTopK(2));String information = similarDocuments.stream().map(Document::getContent).collect(Collectors.joining(System.lineSeparator()));// 构建系统提示模板，包含当前时间和文档内容var systemPromptTemplate =newSystemPromptTemplate("现在的时间是{date}\n"+"你需要使用文档内容对用户提出的问题进行回复，同时你需要表现得天生就知道这些内容，"+"不能在回复中体现出你是根据给出的文档内容进行回复的，这点非常重要。"+"当用户提出的问题无法根据文档内容进行回复或者你也不清楚时，回复不知道即可。"+"文档内容如下：\n"+"{information}");// 将信息和时间格式化为系统消息和用户消息var systemMessage = systemPromptTemplate.createMessage(Map.of("information", information,"date",LocalDateTime.now().format(DateTimeFormatter.ISO_DATE_TIME)));var userMessage =newUserMessage(message);// 将系统消息和用户消息合并，生成回复流return chatModel.stream(newPrompt(List.of(systemMessage, userMessage)).getContents());}

在这个实现中，我们首先使用用户的查询消息对VectorStore进行检索，找到最相关的文档。然后，我们构建一个系统提示模板，将检索到的文档内容和当前时间作为上下文信息。最后，我们将系统提示和用户消息合并，生成一个Prompt对象，并使用OllamaChatModel生成回复流。

结语

通过上述步骤，我们成功地基于Spring-AI框架实现了RAG增强检索。这种方法结合了文档检索和文本生成，能够提供更加丰富和准确的回复。希望本文能够帮助开发者更好地理解和应用Spring-AI框架，开发出更加智能的应用程序。

源码

spring-ai-demo