引言

随着人工智能技术的快速发展，增强检索（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. 准备工作

安装ollama下载大模型

<code>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。

@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框架，开发出更加智能的应用程序。