使用 MongoDB 和 Spring AI 构建 RAG 应用

2024-09-26

1、概览

AI（人工智能）技术的使用正成为现代开发中的一项关键技能。在本文中，我们将构建一个 RAG Wiki 应用，它可以根据存储的文档回答问题。

我们会通过 Spring AI 将应用与 MongoDB Vector 数据库和 LLM 集成。

2、RAG 应用

当自然语言生成需要依赖上下文数据时，我们就会使用 RAG（Retrieval-Augmented Generation）应用。RAG 应用的一个关键组成部分是向量数据库（Vector Database），它在有效管理和检索这些数据方面起着至关重要的作用：

RAG 应用中的文档填充和提示流程

我们使用 Embedding Model 来处理源文件。Embedding Model 将文档中的文本转换为高维向量。这些向量捕捉了内容的语义，使我们能够根据上下文而不仅仅是关键词匹配来比较和检索类似的内容。然后，我们将文档存储到向量数据库。

保存文档后，我们可以通过以下方式发送提示信息：

首先，我们使用 Embedding Model 来处理问题，将其转换为捕捉其语义的向量。
接下来，进行相似性搜索，将问题的向量与存储在向量库中的文档向量进行比较。
从最相关的文档中，为问题建立一个上下文。
最后，将问题及其上下文发送给 LLM，LLM 会根据所提供的上下文构建与查询相关的回复。

3、MongoDB Atlas Search

在本教程中，我们将使用 MongoDB Atlas Search 作为我们的向量数据库。它提供的向量搜索功能可以满足我们在本项目中的需求。为了测试，我们使用 mongodb-atlas-local Docker 容器来设置 MongoDB Atlas Search 的本地实例。

创建一个 docker-compose.yml 文件：

version: '3.1'

services:
  my-mongodb:
    image: mongodb/mongodb-atlas-local:7.0.9
    container_name: my-mongodb
    environment:
      - MONGODB_INITDB_ROOT_USERNAME=wikiuser
      - MONGODB_INITDB_ROOT_PASSWORD=password
    ports:
      - 27017:27017

4、依赖和配置

首先，添加必要的依赖项。由于我们的应用要提供 HTTP API，因此加入 spring-boot-starter-web 依赖：

<dependency>
    <groupId>org.springframework.boot</groupId> 
    <artifactId>spring-boot-starter-web</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

此外，我们使用 Open AI API 客户端连接到 LLM，因此也要添加它的依赖：

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

最后，添加 MongoDB Atlas Store 依赖：

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-mongodb-atlas-store-spring-boot-starter</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

现在，添加配置属性：

spring:
  data:
    mongodb:
      uri: mongodb://wikiuser:password@localhost:27017/admin
      database: wiki
  ai:
    vectorstore:
      mongodb:
        collection-name: vector_store
        initialize-schema: true
        path-name: embedding
        indexName: vector_index
    openai:
      api-key: ${OPENAI_API_KEY}
      chat:
        options:
          model: gpt-3.5-turbo

我们指定了 MongoDB URL 和数据库，还通过设置集合名称、Embedding 字段名称和向量索引名称配置了向量数据库。借助 initialize-schema 属性，Spring AI 框架将自动创建所有这些组件。

最后，我们添加了 Open AI API key 和模型版本。

5、保存文档到向量数据库

现在，我们要实现将数据保存到向量数据库中的过程。我们的应用将负责在现有文档的基础上为用户提供问题解答，本质上就像是一种 Wiki（维基）。

添加一个模型来存储文件内容和文件路径：

public class WikiDocument {
    private String filePath;
    private String content;

    // Getter / Setter 省略
}

下一步是添加 WikiDocumentsRepository。在这个 Repository 中，我们要封装所有的持久化逻辑：

import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;

@Component
public class WikiDocumentsRepository {
    private final VectorStore vectorStore;

    public WikiDocumentsRepository(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    public void saveWikiDocument(WikiDocument wikiDocument) {

        Map<String, Object> metadata = new HashMap<>();
        metadata.put("filePath", wikiDocument.getFilePath());
        Document document = new Document(wikiDocument.getContent(), metadata);
        List<Document> documents = new TokenTextSplitter().apply(List.of(document));

        vectorStore.add(documents);
    }
}

如上，我们注入了 VectorStore 接口 Bean，它的实现是由 spring-ai-mongodb-atlas-store-spring-boot-starter 提供的 MongoDBAtlasVectorStore。在 saveWikiDocument 方法中，我们创建了一个 Document 实例，并在其中填充了内容和元数据。

然后，我们使用 TokenTextSplitter 将文档分割成小块，并保存到向量数据库中。

现在，让我们创建一个 WikiDocumentsServiceImpl：

@Service
public class WikiDocumentsServiceImpl {
    private final WikiDocumentsRepository wikiDocumentsRepository;

    // 构造函数

    public void saveWikiDocument(String filePath) {
        try {
            String content = Files.readString(Path.of(filePath));
            WikiDocument wikiDocument = new WikiDocument();
            wikiDocument.setFilePath(filePath);
            wikiDocument.setContent(content);

            wikiDocumentsRepository.saveWikiDocument(wikiDocument);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
}

在 Service 层，我们检索文件内容，创建 WikiDocument 实例，并将其发送到数据库进行持久化。

在 Controller 中，我们只需将文件路径传递给 Service 层，如果文档保存成功，则返回 201 状态码：

@RestController
@RequestMapping("wiki")
public class WikiDocumentsController {
    private final WikiDocumentsServiceImpl wikiDocumentsService;

    // 构造函数

    @PostMapping
    public ResponseEntity<Void> saveDocument(@RequestParam String filePath) {
        wikiDocumentsService.saveWikiDocument(filePath);

        return ResponseEntity.status(201).build();
    }
}

现在，启动应用，看看流程是如何运行的。

添加 Spring Boot 测试依赖项，它将允许我们设置 Test Web Context：

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <version>LATEST_VERSION</version>
</dependency>

现在，启动测试用例，并调用 POST 端点获取两个文档：

@AutoConfigureMockMvc
@ExtendWith(SpringExtension.class)
@SpringBootTest
class RAGMongoDBApplicationManualTest {

    @Autowired
    private MockMvc mockMvc;

    @Test
    void givenMongoDBVectorStore_whenCallingPostDocumentEndpoint_thenExpectedResponseCodeShouldBeReturned() throws Exception {
        mockMvc.perform(post("/wiki?filePath={filePath}",
          "src/test/resources/documentation/owl-documentation.md"))
          .andExpect(status().isCreated());

        mockMvc.perform(post("/wiki?filePath={filePath}",
          "src/test/resources/documentation/rag-documentation.md"))
          .andExpect(status().isCreated());
    }
}

两次调用都应返回 201 状态码，表明文档已被添加。我们可以使用 MongoDB Compass 来确认文档是否已成功保存到向量数据库中：

通过 MongoDB Compass 查看存储的文档

我们可以看到，两个文档都已保存，可以看到原始内容和 Embedding 数组。

6、相似性搜索

让我们实现相似性搜索功能。

在 Repository 中加入 findSimilarDocuments 方法：

@Component
public class WikiDocumentsRepository {
    private final VectorStore vectorStore;

    public List<WikiDocument> findSimilarDocuments(String searchText) {

        return vectorStore
          .similaritySearch(SearchRequest
            .query(searchText)
            .withSimilarityThreshold(0.87)
            .withTopK(10))
          .stream()
          .map(document -> {
              WikiDocument wikiDocument = new WikiDocument();
              wikiDocument.setFilePath((String) document.getMetadata().get("filePath"));
              wikiDocument.setContent(document.getContent());

              return wikiDocument;
          })
          .toList();
    }
}

我们调用了 VectorStore 中的 similaritySearch 方法。除了搜索文本外，还指定了结果限制和相似度阈值。通过相似度阈值参数，我们可以控制文档内容与搜索文本的匹配程度。

Service 层代理了对 Repository 的调用：

public List<WikiDocument> findSimilarDocuments(String searchText) {
    return wikiDocumentsRepository.findSimilarDocuments(searchText);
}

在 Controller 中，我们添加了一个 GET 端点，接收 searchText 作为参数并将其传递给 Service：

@RestController
@RequestMapping("/wiki")
public class WikiDocumentsController {
    @GetMapping
    public List<WikiDocument> get(@RequestParam("searchText") String searchText) {
        return wikiDocumentsService.findSimilarDocuments(searchText);
    }
}

现在调用新的端点，看看相似性搜索是如何工作的：

@Test
void givenMongoDBVectorStoreWithDocuments_whenMakingSimilaritySearch_thenExpectedDocumentShouldBePresent() throws Exception {
    String responseContent = mockMvc.perform(get("/wiki?searchText={searchText}", "RAG Application"))
      .andExpect(status().isOk())
      .andReturn()
      .getResponse()
      .getContentAsString();

    assertThat(responseContent)
      .contains("RAG AI Application is responsible for storing the documentation");
}

我们使用一个在文档中并非完全匹配的搜索文本调用了端点。然而，我们仍然检索到了具有相似内容的文档，并确认其中包含了我们在 rag-documentation.md 文件中存储的文本。

7、Prompt 端点

实现提示流程，这是应用的核心功能。

从 AdvisorConfiguration 开始：

@Configuration
public class AdvisorConfiguration {

    @Bean
    public QuestionAnswerAdvisor questionAnswerAdvisor(VectorStore vectorStore) {
        return new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults());
    }
}

我们创建了一个 QuestionAnswerAdvisor Bean，负责构建提示请求，包括初始问题。此外，它还会附加向量数据库的相似性搜索响应，作为问题的上下文。

现在，将搜索端点添加到 API 中：

@RestController
@RequestMapping("/wiki")
public class WikiDocumentsController {
    private final WikiDocumentsServiceImpl wikiDocumentsService;
    private final ChatClient chatClient;
    private final QuestionAnswerAdvisor questionAnswerAdvisor;

    public WikiDocumentsController(WikiDocumentsServiceImpl wikiDocumentsService,
                                   @Qualifier("openAiChatModel") ChatModel chatModel,
                                   QuestionAnswerAdvisor questionAnswerAdvisor) {
        this.wikiDocumentsService = wikiDocumentsService;
        this.questionAnswerAdvisor = questionAnswerAdvisor;
        this.chatClient = ChatClient.builder(chatModel).build();
    }

    @GetMapping("/search")
    public String getWikiAnswer(@RequestParam("question") String question) {
        return chatClient.prompt()
          .user(question)
          .advisors(questionAnswerAdvisor)
          .call()
          .content();
    }
}

这里，我们通过将用户输入添加到提示中并附加我们的 QuestionAnswerAdvisor 来构建一个提示请求。

最后，调用端点，看看 RAG 应用告诉我们什么：

@Test
void givenMongoDBVectorStoreWithDocumentsAndLLMClient_whenAskQuestionAboutRAG_thenExpectedResponseShouldBeReturned() throws Exception {
    String responseContent = mockMvc.perform(get("/wiki/search?question={question}", "Explain the RAG Applications"))
      .andExpect(status().isOk())
      .andReturn()
      .getResponse()
      .getContentAsString();

    logger.atInfo().log(responseContent);

    assertThat(responseContent).isNotEmpty();
}

我们向端点发送了 “Explain the RAG applications（解释 RAG 应用）” 的问题，并记录了 API 响应：

b.s.r.m.RAGMongoDBApplicationManualTest : Based on the context provided, the RAG AI Application is a tool 
used for storing documentation and enabling users to search for specific information efficiently...

我们可以看到，端点根据我们之前保存在向量数据库中的文档文件，返回了有关 RAG 应用的信息。

现在，让我们试着问一些知识库中肯定没有的问题：

@Test
void givenMongoDBVectorStoreWithDocumentsAndLLMClient_whenAskUnknownQuestion_thenExpectedResponseShouldBeReturned() throws Exception {
    String responseContent = mockMvc.perform(get("/wiki/search?question={question}", "Explain the Economic theory"))
      .andExpect(status().isOk())
      .andReturn()
      .getResponse()
      .getContentAsString();

    logger.atInfo().log(responseContent);

    assertThat(responseContent).isNotEmpty();
}

我们问到了经济理论，以下是答复：

b.s.r.m.RAGMongoDBApplicationManualTest : I'm sorry, but the economic theory is not directly related to the information provided about owls and the RAG AI Application.
If you have a specific question about economic theory, please feel free to ask.

这一次，我们的应用没有找到任何相关文档，也没有使用其他来源提供答案。

8、总结

在本文中，我们使用 Spring AI 框架成功实现了一个 RAG 应用，该框架是集成各种 AI 技术的绝佳工具。此外，MongoDB 被证明是处理向量存储的最佳选择。

有了这一强大的组合，我们就能为聊天机器人、自动维基（Wiki）系统和搜索引擎等各种用途构建基于 AI 的现代应用。

Ref：https://www.baeldung.com/spring-ai-mongodb-rag