使用向量数据库进行电影相似性搜索
简介
随着 LLM 的普及,我们经常看到它们被用于与文本生成不直接相关的任务。例如,使用 LLM 来构建推荐系统。在这篇文章中,我们将展示如何使用 Quarkus LangChain4j 构建这样一个系统,但无需使用 LLM。更具体地说,我们将使用向量数据库创建一个简单的电影相似性搜索系统。在这个过程中,Quarkus LangChain4j 的作用是通过 `EmbeddingStore` 接口来抽象底层的向量数据库。
最近,一个新的相关示例已添加到 Quarkus LangChain4j 示例中。
嵌入
嵌入是一种将非结构化数据(文本、图像等)结构化的方式。这通过将数据映射到向量来实现。由于我们可以对向量执行数学运算,例如计算它们之间的距离,因此我们可以使用嵌入来计算两块数据的接近程度(或相似度)。在我们的案例中,比较电影概述的嵌入可以让我们衡量两部电影的相似程度。这就是本文的前提。
如何创建嵌入?
Quarkus LangChain4j 提供了以下几种创建嵌入的方法:
在这篇文章中,我们将使用前者。我有没有提到我们不会使用 LLM?相反,我们将使用 pgvector,这是一个 PostgreSQL 扩展,提供向量操作和索引。
<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-embeddings-bge-small-en-q</artifactId>
</dependency>
<dependency>
<groupId>io.quarkiverse.langchain4j</groupId>
<artifactId>quarkus-langchain4j-pgvector</artifactId>
</dependency>
为了能够使用这些依赖项而无需指定版本,可以将 BOM 导入项目的 `dependencyManagement` 中。
<dependency>
<groupId>io.quarkiverse.langchain4j</groupId>
<artifactId>quarkus-langchain4j-bom</artifactId>
<version>0.25.0</version>
<type>pom</type>
<scope>import</scope>
</dependency>
为了正确使用进程内嵌入模型,我们需要在 `application.properties` 文件中对其进行配置。我们还需要配置 pgvector 的维度,并确保它与嵌入模型的维度保持一致。在我们的案例中,它是 384(Quarkus LangChain4j 文档提供了每个模型使用的尺寸)。
因此,`application.properties` 文件应如下所示:
quarkus.langchain4j.pgvector.dimension=384
quarkus.langchain4j.embedding-model.provider=dev.langchain4j.model.embedding.onnx.bgesmallenq.BgeSmallEnQuantizedEmbeddingModel
注意:我们可以使用 Quarkus LangChain4j 支持的任何其他文档存储,事实上,这是使用它的关键优势之一,即对嵌入存储的抽象。
存储嵌入
为了存储嵌入,我们需要一个 `EmbeddingStoreIngestor`。Ingestor 是使用 `EmbeddingModel` 和 `EmbeddingStore` 创建的,两者都作为 bean 提供,并且可以轻松注入。或者,用户可以指定一个文档拆分器,用于将大型文档拆分成较小的块,但在这篇文章中我们不需要它,因为电影概述相对较小。

在将非结构化文本传递给 ingestor 之前,我们需要将其包装在 `Document` 对象中。`Document` 还包含一个 `Metadata` 对象,其中保存键值对。`Metadata` 非常有用,因为我们可以添加需要将 `Document` 与其他数据相关联的信息。在我们的案例中,我们将使用 `Metadata` 来存储电影的 ID。该 ID 将帮助我们关联电影概述和实际电影。
下面的简化代码展示了实际的 ingestor 是如何创建的,以及嵌入是如何存储的。
@Inject
EmbeddingModel embeddingModel;
@Inject
EmbeddingStore embeddingStore;
public EmbeddingStoreIngestor createIngestor() {
return EmbeddingStoreIngestor.builder()
.embeddingModel(embeddingModel)
.embeddingStore(embeddingStore)
.build();
}
public void ingest(Long movieId, String overview) {
Metadata metadata = Metadata.from(Map.of("id", id));
Document document = Document.from(overview, metadata);
createIngestor().ingest(document);
}
那么,我们到底如何使用电影 ID 呢?这实际上取决于我们如何存储其余的电影数据。在我们的案例中,我们将电影数据存储在 PostgreSQL 数据库中。这意味着电影 ID 对应于数据库中电影的 ID。
查询嵌入
为了查询嵌入,我们将使用 `EmbeddingStore` 和 `EmbeddingModel` 来执行 `EmbeddingSearchRequest`。代码非常直接。我们使用电影概述来创建搜索请求。`EmbeddingSearchRequest` 构建器还允许我们指定最大结果数和最小相似度阈值。后者允许我们过滤掉与查询嵌入不够相似的嵌入。换句话说,它告诉存储,如果没有足够相似的结果,请避免在响应中添加不相关的结果。
package io.quarkiverse.langchain4j.sample;
import java.util.List;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.inject.Inject;
import jakarta.transaction.Transactional;
@ApplicationScoped
public class MovieRecommendationService {
@Inject
EmbeddingStore<TextSegment> embeddingStore;
@Inject
EmbeddingModel embeddingModel;
@Transactional
public List<Movie> searchSimilarMovies(String overview) {
Embedding embedding = embeddingModel.embed(overview).content();
EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
.queryEmbedding(embedding)
.minScore(0.5)
.maxResults(10)
.build();
return embeddingStore.search(request).matches().stream().map(m -> {
Long id = m.embedded().metadata().getLong("id");
Movie movie = Movie.findById(id);
return movie;
}).toList();
}
}
加载电影
为了将电影填充到数据库中,我们将使用一个 CSV 文件,其中包含 IMDB 排名前 1000 的电影。对我们而言,重要的列是:
-
`title`:电影标题
-
`overview`:电影概述
-
`link`:海报图片的链接
为了方便地将 CSV 条目映射到 `Movie` 对象,我们将使用 `Jackson`。具体来说,我们将使用 `@JsonProperty` 注释将 CSV 列映射到 `Movie` 字段。此外,我们将使用 `@JsonIgnoreProperties(ignoreUnknown = true)` 来忽略未知字段。
因此,我们的 Movie 实体对象的简化版本如下所示:
package org.acme;
import io.quarkus.hibernate.orm.panache.PanacheEntity;
import jakarta.persistence.Column;
import jakarta.persistence.Entity;
import java.util.List;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
@Entity
@JsonIgnoreProperties(ignoreUnknown = true)
public class Movie extends PanacheEntity {
@JsonProperty("Poster_Link")
public String link;
@JsonProperty("Series_Title")
public String title;
@JsonProperty("Overview")
@Column(length = 1000)
public String overview;
public static List<Movie> searchByTitleLike(String title) {
return find("title like ?1", "%" + title + "%").list();
}
}
现在,我们准备好将电影从 CSV 加载到我们的关系数据库和向量数据库中了。
要点
-
观察 `StartupEvent` 允许我们在应用程序启动时加载电影。
-
CSVMapper 用于将 CSV 条目映射到 `Movie` 对象。
-
我们使用 `@Transaction` 方法保存每部电影,因为我们需要通常由数据库生成的 `id`。
-
我们分批摄取文档。这可以产生巨大的差异,尤其是在我们从进程内嵌入模型迁移到远程嵌入模型时。
package org.acme;
import io.quarkus.logging.Log;
import io.quarkus.runtime.StartupEvent;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.event.Observes;
import jakarta.transaction.Transactional;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.eclipse.microprofile.config.inject.ConfigProperty;
import com.fasterxml.jackson.databind.MappingIterator;
import com.fasterxml.jackson.dataformat.csv.CsvMapper;
import com.fasterxml.jackson.dataformat.csv.CsvSchema;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
@ApplicationScoped
public class MovieLoader {
public void load(@Observes StartupEvent event, @ConfigProperty(name = "movies.file") Path moviesFile,
EmbeddingStore embeddingStore, EmbeddingModel embeddingModel) throws Exception {
if (!Files.exists(moviesFile)) {
throw new IllegalStateException("Missing movies file: " + moviesFile);
}
embeddingStore.removeAll();
EmbeddingStoreIngestor ingester = EmbeddingStoreIngestor.builder()
.embeddingModel(embeddingModel)
.embeddingStore(embeddingStore)
.build();
List<Document> docs = new ArrayList<>();
try (MappingIterator<Movie> it = new CsvMapper().readerFor(Movie.class).with(CsvSchema.emptySchema().withHeader()).readValues(moviesFile.toFile())) {
for (Movie movie : it.readAll()) {
Long id = save(movie).id;
Metadata metadata = Metadata.from(Map.of("id", id, "title", movie.title));
Document document = Document.from(movie.overview, metadata);
docs.add(document);
}
}
Log.info("Ingesting movies...");
ingester.ingest(docs);
Log.info("Application initalized!");
}
@Transactional
public Movie save(Movie m) {
m.persist();
return m;
}
}
要使用 CSV 映射器,我们需要 Jackson 的 CSV 数据格式依赖项。
<dependency>
<groupId>com.fasterxml.jackson.dataformat</groupId>
<artifactId>jackson-dataformat-csv</artifactId>
</dependency>
电影文件的路径在 `application.properties` 文件中指定。
movies.file=src/main/resources/movies.csv
整合
剩下的就是创建一个 REST 端点,允许我们搜索相似的电影。我们也可以使用一个简单的 UI。让我们先从 REST 端点开始。这很简单。我们需要两个方法,一个用于搜索电影,一个用于搜索相似电影。前者我们只使用 `Movie` 实体,后者我们注入并使用我们之前创建的 `MovieRecommendationService`。
package io.quarkiverse.langchain4j.sample;
import jakarta.inject.Inject;
import jakarta.ws.rs.GET;
import jakarta.ws.rs.Path;
import jakarta.ws.rs.Produces;
import jakarta.ws.rs.core.MediaType;
import java.util.List;
@Path("/movies")
public class MovieResource {
@Inject
MovieRecommendationService recommendationService;
@GET
@Produces(MediaType.APPLICATION_JSON)
@Path("/by-title/{title}")
public List<Movie> searchByTitle(String title) {
return Movie.searchByTitleLike(title);
}
@GET
@Produces(MediaType.APPLICATION_JSON)
@Path("/similar/{id}")
public List<Movie> searchSimilar(Long id) {
Movie m = Movie.findById(id);
return recommendationService.searchSimilarMovies(m.overview);
}
}
对于 UI,我们只需要使用一个简单的 HTML 页面,该页面使用 REST 端点来搜索相似的电影。
该页面的关键元素是:
-
movie-box:用于输入电影标题的文本字段
-
search-results:用于显示搜索结果的无序列表
-
movie-overview:用于显示所选电影概述的 div
-
movie-poster:用于显示电影海报的图像
-
similar-results:用于显示相似电影的额外无序列表
重要的是要记住,`Movie` 实体使用 Jackson 将 CSV 列映射到实体字段。这意味着当 `Movie` 序列化为 JSON 时,它将使用 CSV 列名作为字段名,而不是实体字段名。下面的 HTML 代码需要考虑这一点。
<!DOCTYPE html>
<html lang="en">
<head>
<title>{{_title_}}</title>
<meta charset="UTF-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
</head>
<body>
<h2>Movie Similarity Search</h2>
<input type="text" id="movie-box" placeholder="Enter a movie title">
<h3 hidden="true" id="movie-results-heading">Click on of the movies below</h3>
<ul id="search-results"></ul>
<img id="movie-poster"><img>
<div id="movie-overview"></div>
<h3 hidden="true" id="similar-heading">Similar movies</h3>
<ul id="similar-results"></ul>
<script>
document.getElementById("movie-box").addEventListener("input", async function() {
const query = this.value.trim();
if (query.length === 0) {
document.getElementById("search-results").innerHTML = "";
return;
}
const response = await fetch(`/movies/by-title/${encodeURIComponent(query)}`);
const movies = await response.json();
if (movies.length > 0) {
document.getElementById("movie-results-heading").hidden = false;
}
movies.forEach(movie => {
const li = document.createElement("li");
li.textContent = movie.Series_Title;
li.addEventListener("click", () => displayMovie(movie));
document.getElementById("search-results").appendChild(li);
});
});
async function displayMovie(movie) {
console.log('Displaying movie:', movie);
document.getElementById("search-results").innerHTML = "";
document.getElementById("movie-poster").src = movie.Poster_Link;
document.getElementById("movie-poster").style.display = "block";
document.getElementById("movie-overview").textContent = movie.Overview;
document.getElementById("similar-heading").hidden = false;
document.getElementById("movie-results-heading").hidden = true;
document.getElementById("similar-results").innerHTML = "";
const response = await fetch(`/movies/similar/${encodeURIComponent(movie.id)}`);
const similarMovies = await response.json();
similarMovies.forEach(similarMovie => {
const li = document.createElement("li");
li.textContent = similarMovie.Series_Title;
li.addEventListener("click", () => displayMovie(similarMovie));
document.getElementById("similar-results").appendChild(li);
});
}
</script>
</body>
</html>
我不会详细介绍 HTML 代码,因为它超出了本文的范围。最终结果应如下所示:
