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@ServiceConnectionpublic ChromaDBContainer chromaDB() {return new ChromaDBContainer("chromadb/chroma:0.5.20");}@Bean@ServiceConnectionpublic 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: trueollama:embedding:options:model: nomic-embed-textinit:chat:include: falsepull-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@Overridepublic 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 会先将查询内容转换为向量表示，再去向量存储中执行搜索。

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

---

6. 总结

本文介绍了如何将 ChromaDB 向量存储与 Spring AI 集成。

我们通过Testcontainers启动了ChromaDB 和 Ollama 的 Docker 容器，搭建了本地测试环境。演示了如何在应用启动时通过 PoetryDB API 向向量存储中填充诗歌数据。最后，我们用常见的诗歌主题验证了语义搜索功能。

私信1v1直连大厂总监·「免。米」