Spring AI系列之Spring AI 集成 ChromaDB 向量数据库

1. 概述

在传统数据库中，我们通常依赖精确的关键词或基本的模式匹配来实现搜索功能。虽然这种方法对于简单的应用程序已经足够，但它无法真正理解自然语言查询背后的含义和上下文。

向量存储解决了这一限制，它通过将数据以数值向量的形式存储，从而捕捉数据的语义。相似的词会在向量空间中靠得很近，这使得语义搜索成为可能——即使查询中没有包含确切的关键词，也能返回相关结果。

在本教程中，我们将探讨如何将 ChromaDB（一款开源的向量存储库）与 Spring AI 集成使用。

为了将文本数据转换成 ChromaDB 能够存储和搜索的向量，我们需要使用一个嵌入模型。我们将使用 Ollama 在本地运行嵌入模型。

2. 依赖项

首先，我们需要在项目的 pom.xml 文件中添加必要的依赖项：

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-chroma-store-spring-boot-starter</artifactId>
    <version>1.0.0-M6</version>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-ollama-spring-boot-starter</artifactId>
    <version>1.0.0-M6</version>
</dependency>

ChromaDB starter依赖项使我们能够建立与 ChromaDB 向量存储的连接并与之交互。

此外，我们还引入了 Ollama 启动器依赖项，用于运行我们的嵌入模型。

由于当前版本（1.0.0-M6）是一个里程碑版本，我们还需要在 pom.xml 文件中添加 Spring Milestones 仓库：

<repositories>
    <repository>
        <id>spring-milestones</id>
        <name>Spring Milestones</name>
        <url>https://repo.spring.io/milestone</url>
        <snapshots>
            <enabled>false</enabled>
        </snapshots>
    </repository>
</repositories>

这个仓库用于发布里程碑版本，不同于标准的 Maven Central 仓库。

由于我们在项目中使用了多个SpringAI启动器(starter)，我们还需要在pom.xml中引入SpringAI的BOM(依赖管理清单，Bill of Materials)：

<dependencyManagement>
    <dependencies>
        <dependency>
            <groupId>org.springframework.ai</groupId>
            <artifactId>spring-ai-bom</artifactId>
            <version>1.0.0-M6</version>
            <type>pom</type>
            <scope>import</scope>
        </dependency>
    </dependencies>
</dependencyManagement>

通过引入 BOM，我们现在可以从两个启动器依赖中移除版本号标签。

BOM 的作用是消除版本冲突的风险，并确保我们使用的 Spring AI 相关依赖彼此兼容。

3. 使用 Testcontainers 搭建本地测试环境

为了便于本地开发和测试，我们将使用Testcontainers来搭建 ChromaDB 向量存储和 Ollama 服务。

通过Testcontainers运行这些服务的前提是本机已安装并正在运行 Docker。

3.1 测试依赖项

首先，我们需要在 pom.xml 中添加必要的测试依赖项：

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-spring-boot-testcontainers</artifactId>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.testcontainers</groupId>
    <artifactId>chromadb</artifactId>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.testcontainers</groupId>
    <artifactId>ollama</artifactId>
    <scope>test</scope>
</dependency>

这些依赖项为我们提供了必要的类，可以快速创建用于两个外部服务的临时 Docker 实例。

3.2 定义 Testcontainers Bean

接下来，我们创建一个带有

@TestConfiguration 注解的类，用于定义我们的 Testcontainers Bean：

@TestConfiguration(proxyBeanMethods = false)
class TestcontainersConfiguration {

    @Bean
    @ServiceConnection
    public ChromaDBContainer chromaDB() {
        return new ChromaDBContainer("chromadb/chroma:0.5.20");
    }

    @Bean
    @ServiceConnection
    public OllamaContainer ollama() {
        return new OllamaContainer("ollama/ollama:0.4.5");
    }
}

我们为容器指定了最新的稳定版本。

同时，我们使用 @ServiceConnection 注解标注了 Bean 方法。该注解会自动注册所有连接外部服务所需的属性，从而建立连接。

即使不使用 Testcontainers 支持，Spring AI 在本地运行时也会自动连接到ChromaDB和Ollama，前提是它们分别运行在默认端口 8000 和 11434 上。

然而，在生产环境中，我们可以通过相应的 Spring AI 配置属性来覆盖这些连接信息：

spring:
  ai:
    vectorstore:
      chroma:
        client:
          host: ${CHROMADB_HOST}
          port: ${CHROMADB_PORT}
    ollama:
      base-url: ${OLLAMA_BASE_URL}

一旦连接信息配置正确，SpringAI会自动为我们创建 VectorStore

和 EmbeddingModel 类型的 Bean，

使我们可以分别与向量存储和嵌入模型进行交互。我们将在本教程后续部分了解如何使用这些 Bean。

尽管@ServiceConnection会自动定义所需的连接信息，但我们仍然需要在 application.yml 文件中配置一些额外的属性：

spring:
  ai:
    vectorstore:
      chroma:
        initialize-schema: true
    ollama:
      embedding:
        options:
          model: nomic-embed-text
      init:
        chat:
          include: false
        pull-model-strategy: WHEN_MISSING

在这里，我们启用了ChromaDB的schema初始化。接着，我们将nomic-embed-text配置为嵌入模型，并指示Ollama在本地未安装该模型时自动拉取。

当然，我们也可以根据需求使用Ollama中的其他嵌入模型，或是 Hugging Face 的模型。

3.3 开发过程中使用 Testcontainers

虽然 Testcontainers 主要用于集成测试，但我们也可以在本地开发时使用它。

为此，我们将在src/test/java目录中创建一个单独的主类(main class)：

class TestApplication {

    public static void main(String[] args) {
        SpringApplication.from(Application::main)
          .with(TestcontainersConfiguration.class)
          .run(args);
    }
}

我们创建了一个TestApplication类，

并在其main方法中，通过

TestcontainersConfiguration 启动我们的主应用类 Application。

这个配置有助于我们在本地搭建和管理外部服务。我们可以运行Spring Boot 应用程序，并让它连接通过 Testcontainers 启动的外部服务。

4. 在应用启动时填充 ChromaDB

现在我们已经搭建好了本地环境，接下来在应用启动时向 ChromaDB 向量存储中填充一些示例数据。

4.1 从 PoetryDB 获取诗歌数据

为了演示，我们将使用 PoetryDB API 获取一些诗歌。

我们先创建一个名为 PoetryFetcher 的工具类来实现这个功能：

class PoetryFetcher {

    private static final String BASE_URL = "https://poetrydb.org/author/";
    private static final String DEFAULT_AUTHOR_NAME = "Shakespeare";

    public static List<Poem> fetch() {
        return fetch(DEFAULT_AUTHOR_NAME);
    }

    public static List<Poem> fetch(String authorName) {
        return RestClient
          .create()
          .get()
          .uri(URI.create(BASE_URL + authorName))
          .retrieve()
          .body(new ParameterizedTypeReference<>() {});
    }

}

record Poem(String title, List<String> lines) {}

我们使用 RestClient 调用 PoetryDB API，并传入指定的 authorName。为了将API响应反序列化为Poem记录的列表，我们使用了 ParameterizedTypeReference，无需显式指定泛型类型，Java 会为我们自动推断类型。我们还重载了fetch()方法(不带任何参数)，以便默认获取莎士比亚（Shakespeare）的诗歌。

4.2 将文档存入 ChromaDB 向量存储

现在，为了在应用启动时将诗歌填充到 ChromaDB 向量存储中，我们将创建一个 VectorStoreInitializer 类，并实现 ApplicationRunner 接口：

@Component
class VectorStoreInitializer implements ApplicationRunner {

    private final VectorStore vectorStore;

    // standard constructor

    @Override
    public void run(ApplicationArguments args) {
        List<Document> documents = PoetryFetcher
          .fetch()
          .stream()
          .map(poem -> {
              Map<String, Object> metadata = Map.of("title", poem.title());
              String content = String.join("\n", poem.lines());
              return new Document(content, metadata);
          })
          .toList();
        vectorStore.add(documents);
    }

}

在我们的VectorStoreInitializer中，我们通过自动注入（@Autowired）获取了一个 VectorStore 实例。

在run()方法中，我们使用

PoetryFetcher 工具类获取诗歌列表。

然后，将每首诗映射成一Document，将诗的内容(lines)作为文本内容，标题(title)作为元数据。

最后，我们将所有文档存入向量存储。当调用 add() 方法时，Spring AI 会自动将纯文本内容转换成向量表示，然后存入向量存储中，无需我们显式调用 EmbeddingModel Bean 进行转换。

默认情况下，Spring AI使用

SpringAiCollection作为向量存储中

的集合名称，但我们可以通过spring.ai.vectorstore.chroma.collection-name

属性来覆盖该名称。

5. 测试语义搜索

在 ChromaDB 向量存储被填充数据后，让我们来验证语义搜索功能：

private static final int MAX_RESULTS = 3;

@ParameterizedTest
@ValueSource(strings = {"Love and Romance", "Time and Mortality", "Jealousy and Betrayal"})
void whenSearchingShakespeareTheme_thenRelevantPoemsReturned(String theme) {
    SearchRequest searchRequest = SearchRequest
      .builder()
      .query(theme)
      .topK(MAX_RESULTS)
      .build();
    List<Document> documents = vectorStore.similaritySearch(searchRequest);

    assertThat(documents)
      .hasSizeLessThanOrEqualTo(MAX_RESULTS)
      .allSatisfy(document -> {
          String title = String.valueOf(document.getMetadata().get("title"));
          assertThat(title)
            .isNotBlank();
        });
}

这里，我们使用 @ValueSource 向测试方法传入一些常见的莎士比亚主题。然后，创建一个SearchRequest对象，将主题作为查询内容，将 MAX_RESULTS 设为期望的返回结果数量。

接着，我们调用 vectorStore Bean 的 similaritySearch() 方法，传入我们的 searchRequest。与 add() 方法类似，Spring AI 会先将查询内容转换为向量表示，再去向量存储中执行搜索。

返回的文档中包含的诗歌都是与给定主题语义相关的，即使它们没有包含确切的关键词。