基于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>
    
    <dependency>
        <groupId>org.springframework.bootgroupId>
        <artifactId>spring-boot-starter-webfluxartifactId>
    dependency>
    
    <dependency>
        <groupId>org.springframework.sessiongroupId>
        <artifactId>spring-session-coreartifactId>
    dependency>
    
    <dependency>
        <groupId>org.springframework.aigroupId>
        <artifactId>spring-ai-ollamaartifactId>
    dependency>
    
    <dependencyManagement>
        <dependencies>
            <dependency>
                <groupId>org.springframework.aigroupId>
                <artifactId>spring-ai-bomartifactId>
                <version>1.0.0-SNAPSHOTversion>
                <type>pomtype>
                <scope>importscope>
            dependency>
        dependencies>
    dependencyManagement>
dependencies>

4. OllamaAutoConfig配置

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

@Configuration
public class OllamaAutoConfig {
    private static final Logger log = LoggerFactory.getLogger(OllamaAutoConfig.class);

    @Bean
    public OllamaApi ollamaApi(){
        return new OllamaApi("http://localhost:11434");
    }
    @Bean
    public OllamaChatModel chatModel(OllamaApi ollamaApi) {
        OllamaChatModel chatModel = new OllamaChatModel(ollamaApi,
                OllamaOptions.create()
                        .withModel("wangshenzhi/llama3-8b-chinese-chat-ollama-q4")
                        .withTemperature(0.9f));
        return chatModel;
    }
    @Bean
    public OllamaEmbeddingModel embeddingModel(OllamaApi ollamaApi){
        return new OllamaEmbeddingModel(ollamaApi).withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
    }
    @Bean
    public VectorStore vectorStore(OllamaEmbeddingModel embeddingModel) {
        log.info("开始加载本地知识库文档");
        SimpleVectorStore vectorStore = new SimpleVectorStore(embeddingModel);
        File files = new File(ClassLoader.getSystemResource("doc").getFile());

        for(File f:files.listFiles()){
            Resource resource = new FileSystemResource(f);
            TextReader textReader = new TextReader(resource);
            textReader.setCharset(Charset.defaultCharset());
            vectorStore.add(textReader.get());
            log.info("知识库文件加载完成：{}",f.getName());
        }
        return vectorStore;
    }
}

5. Session配置

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

@Configuration
@EnableSpringWebSession
public class SessionConfig {
    @Bean
    public ReactiveSessionRepository reactiveSessionRepository() {
        return new ReactiveMapSessionRepository(new ConcurrentHashMap<>());
    }
}

6. AI对话接口

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

@GetMapping("/ai/generateStream")
    public Flux<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 = new Prompt(message);
            }else {
                List<Message> messages = new ArrayList<>();
                messages.add(new UserMessage(prompt.getContents()));
                messages.add(new UserMessage(message));
                prompt = new Prompt(messages);
            }
            webSession.getAttributes().put("prompt",prompt);
            return Mono.just(prompt);
        });
        Flux<String> stream = promptMono.flatMapMany(this::getReply);

        return stream;
    }
    private Flux<String> getReply(Prompt prompt){
        return chatModel.stream(prompt.getContents());
    }

7. 实现RAG增强检索

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

@GetMapping("/ai/query")
public Flux<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 = new SystemPromptTemplate(
            "现在的时间是{date}\n" +
            "你需要使用文档内容对用户提出的问题进行回复，同时你需要表现得天生就知道这些内容，" +
            "不能在回复中体现出你是根据给出的文档内容进行回复的，这点非常重要。" +
            "当用户提出的问题无法根据文档内容进行回复或者你也不清楚时，回复不知道即可。" +
            "文档内容如下：\n" +
            "{information}");

    // 将信息和时间格式化为系统消息和用户消息
    var systemMessage = systemPromptTemplate.createMessage(
            Map.of("information", information, "date", LocalDateTime.now().format(DateTimeFormatter.ISO_DATE_TIME)));
    var userMessage = new UserMessage(message);

    // 将系统消息和用户消息合并，生成回复流
    return chatModel.stream(new Prompt(List.of(systemMessage, userMessage)).getContents());
}

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

结语

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