Spring AI 支持向量数据库源码分析与PGVector集成示例

1,248 阅读8分钟

# 12. Spring AI Embedding模型概念、源码分析和使用示例 介绍了嵌入模型,嵌入就是将高维度的数据转换为向量,那么向量应该如何存储呢?Spring AI框架都集成了哪些向量数据库。本文将详细介绍向量数据库的使用。

了解向量是什么?

理解向量相关的知识是理解大模型的底层算法的关键,也是理解向量数据库的关键。简单的复习一下数学知识。 image.png 向量是有维度(或者大小)和方向的。既然有大小和方向,所以向量是可以计算的,比如向量相加、向量相乘以及向量的相似度等。「向量数据库就是使用到了向量的相似原理」

向量相加

向量相加:两条向量一共走了多少路径的直线距离。

简单生活中的例子:从南京去北京,A 直接坐飞机直线,B坐高铁经过各城市 最终两人距离是一样的。

A:南京 ---> 北京

B:南京 ---> 徐州 ---> 济南 ---> 天津 ---> 北京 (中间过程可以理解为向量)

A = B, 我们可以理解为向量相加

向量相乘

向量相乘:一个向量在另一个向量的正射投影与该向量长度的乘积。

简单生活中的例子:两个人拉车子,A直着接拉,B与A有一定的角度拉,我们算一下两个人一共拉车的力量有多大

向量相似

一图形象的描述两个向量是相似: image.png 向量相似

  • 余弦相似度(Cosine Similarity):用于度量两个向量之间的夹角余弦值,取值范围为[-1, 1],其中1表示完全相似,-1表示完全不相似。

  • 皮尔逊相关系数(Pearson Correlation Coefficient):用于度量两个变量之间的线性相关性,取值范围[-1, 1],其中1表示完全正相关,-1表示完全负相关,0表示不相关。

  • 欧氏距离(Euclidean Distance):用于度量两个向量之间的欧氏距离,即两个向量各个元素差的平方和再开根号,取值范围为[0, +∞),其中0表示完全相同。

  • 曼哈顿距离(Manhattan Distance):用于度量两个向量之间的曼哈顿距离,即两个向量各个元素差的绝对值之和,取值范围为[0, +∞) ,其中0表示完全相同。

Q:为什么我们可以在生活中区分不同的物品和事物?比如看到一个动物,我们可以分辨出是猫还是狗,看到一个人可以分辨其是女性还是男性。

A: 我们能够识别物品和食物的特征,比如不同种类的小狗,就可以通过体型大小、毛发长度、鼻子长短等特征来区分。

向量数据库能干什么?

向量数据库的核心在于相似性搜索(Similarity Search),这是向量数据库区别于传统数据库的重要特性。那么如何实现的?

我们通过识别不同事物之间不同的特征来识别种类,例如分别不同种类的小狗,就可以通过体型大小、毛发长度、鼻子长短等特征来区分。对于这些特征我们可以转化为向量,比如按照体型大小,这样就能得到一个体型特征的一维坐标和对应的数值,从 0 到 1 的数字中得到每只狗在坐标系中的位置。 image.png

比如狗还可以从很多种维度进行向量化,[体型大小、毛发长度、鼻子长短] -> [xxx,xxx,xxxx] 的向量。实际上,只要维度够多,我们就能够将所有的事物区分开来,世间万物都可以用一个多维坐标系来表示,它们都在一个高维的特征空间中对应着一个坐标点。

如何将生成和挑选特征这个过程?Feature Engineering (特征工程),它是将原始数据转化成更好的表达问题本质的特征的过程。

向量数据库的作用就是对向量进行存储,并进行相似性搜索。

框架支持哪些向量数据库?

向量数据库官方地址
Azure Vector Search The Azure vector store.
Apache CassandraThe Apache Cassandra vector store.
Chroma Vector Store The Chroma vector store.
Elasticsearch Vector StoreThe Elasticsearch vector store.
GemFire Vector Store The GemFire vector store.
Milvus Vector Store The Milvus vector store.
MongoDB Atlas Vector StoreThe MongoDB Atlas vector store.
Neo4j Vector StoreThe Neo4j vector store.
OpenSearch Vector StoreThe OpenSearch vector store.
Oracle Vector StoreThe Oracle Database vector store.
PgVector StoreThe PostgreSQL/PGVector vector store.
Pinecone Vector Store PineCone vector store.
Qdrant Vector Store Qdrant vector store.
Redis Vector StoreThe Redis vector store.
SAP Hana Vector Store The SAP HANA vector store.
Typesense Vector Store The Typesense vector store.
Weaviate Vector Store The Weaviate vector store.
SimpleVectorStore用于学习目的,简单的向量存储引擎

安装PGVector

  • Docker上安装
docker pull pgvector/pgvector:pg16

docker run --name pgvector --restart=always -e POSTGRES_USER=pgvector -e POSTGRES_PASSWORD=pgvector -p 5432:5432 -d pgvector/pgvector:pg16
  • Linux and Mac上安装
cd /tmp
git clone --branch v0.7.2 https://github.com/pgvector/pgvector.git
cd pgvector
make
make install # may need sudo

安装完成后,可以Spring AI官网说明文档进行操作一下。更详细内容可以参考

对于PgVector的使用大家可参考官方文档,比如如何建表、如何插入数据、如何查询、如何创建索引等。

源码分析

VectorStore 接口

public interface VectorStore extends DocumentWriter {

    // 向向量数据库中添加一组Document对象
    void add(List<Document> documents);

    @Override
    default void accept(List<Document> documents) {
       add(documents);
    }

    // 根据文档ID列表,删除向量数据库的中Docuemnt
    Optional<Boolean> delete(List<String> idList);

    // 根据请求从向量数据库搜索相似的文档,在request中设置一些参数
    List<Document> similaritySearch(SearchRequest request);

    // 根据查询文本进行相似文档的搜索,使用默认的搜索规则
    default List<Document> similaritySearch(String query) {
       return this.similaritySearch(SearchRequest.query(query));
    }
}

SearchRequest

在查询时,根据一些业务场景可以指定一些返回控制参数,比如返回topK个文档,指定相似度或者过滤条件等。

public class SearchRequest {
   // 查询文本
	public final String query;
   // 指定查询返回最相似的前K个,默认值为4,我们可以在查询前指定
	private int topK = DEFAULT_TOP_K;
   // 查询返回相似度的阙值,只有大于该阙值的文档才返回,默认为0.0
	private double similarityThreshold = SIMILARITY_THRESHOLD_ALL;
   // 原数据过滤,在进行嵌入时,可以在原数据中放入一些kv,在查询的时候可以过滤
   // 指定该字段可以过滤,比如权限控制等。
	private Filter.Expression filterExpression;

	public static SearchRequest query(String query) { return new SearchRequest(query); }

	private SearchRequest(String query) { this.query = query; }

	public SearchRequest withTopK(int topK) {...}
	public SearchRequest withSimilarityThreshold(double threshold) {...}
	public SearchRequest withSimilarityThresholdAll() {...}
	public SearchRequest withFilterExpression(Filter.Expression expression) {...}
	public SearchRequest withFilterExpression(String textExpression) {...}

	public String getQuery() {...}
	public int getTopK() {...}
	public double getSimilarityThreshold() {...}
	public Filter.Expression getFilterExpression() {...}
}

Filter.Expression底层使用antlr4技术实现,antlr4主要是词法分析框架。

antlr4 官网:www.antlr.org/download.ht…

VectorStore 实现类

image.png

特别提醒:在选择嵌入模型和大模型时,两者对嵌入向量的维度要匹配

假如:大模型选择的gemma:2b,向量数据选择的PgVector 默认情况下会出现 ERROR: expected 1536 dimensions, not 2048 错误。因为PgVector默认1536, gemma:2b返回2048维度。

代码实战

# 12. Spring AI Embedding模型概念、源码分析和使用示例基础上进行。

applicaiton.yml配置

spring:
  datasource:
    url: jdbc:postgresql://localhost:5432/postgres
    username: pgvector
    password: pgvector
  ai:
    ollama:
      base-url: http://localhost:11434
      embedding:
        options:
          model: gemma:2b
    vectorstore:
      pgvector:
        index-type: hnsw # 索引类型
        distance-type: cosine_distance # 指定相似度计算方式

通过测试,发现Spring AI框架对集成PgVector有个Bug,如果index-type: hnsw/ivfflat时,就会在embedding列上创建索引,我们使用的大模型gemma:2b返回的2048维度,此时启动程序就会报错误,原因如下;

PgVector 创建索引的列的维度最大为2000,如果返回2048 当然会报错。

在使用时,一定要注意了。有Bug的代码如下; image.png

代码

插入一组Document

@Test
public void testAdd() {
    List<Document> documents = List.of(
            new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("year", "2021", "user", "zhang")),
            new Document("The World is Big and Salvation Lurks Around the Corner", Map.of("year", "2019", "user", "li")),
            new Document("You walk forward facing the past and you turn back toward the future.", Map.of("year", "2024", "user", "wang")));

    vectorStore.accept(documents);
}

测试结果如下,默认的表结构就四个字段。 image.png

相似性查询

/**
 * 测试相似查询,仅设置query参数
 */
@Test
public void testSimilaritySearch() {
    List<Document> documents = vectorStore.similaritySearch("Spring AI rocks!!");
    Assertions.assertEquals(documents.size(), 3);
}

/**
 * 相似度查询,设置相似度阈值和返回结果数
 */
@Test
public void testSimilaritySearchWithRequest() {
    List<Document> documents = vectorStore.similaritySearch(
            SearchRequest
                    .query("Spring AI rocks!!")
                    .withTopK(2)
                    .withSimilarityThreshold(0.8));

    Assertions.assertEquals(documents.size(), 1);
}

相似查询 + Filter

/**
 * 查询相似度,并且是zhang的数据
 */
@Test
public void testSimilaritySearchWithRequestAndFilter() {
    //  FilterExpressionBuilder方式,这种方式更容易使用一些
    FilterExpressionBuilder b = new FilterExpressionBuilder();
    List<Document> documents = vectorStore.similaritySearch(
            SearchRequest
                    .query("Spring AI rocks!!")
                    .withFilterExpression(b.eq("user", "zhang").build())
    );

    Assertions.assertEquals(documents.size(), 1);

    // 另外一种写法
    List<Document> docs = vectorStore.similaritySearch(
            SearchRequest
                    .query("Spring AI rocks!!")
                    .withFilterExpression("user == 'zhang'")
    );

    Assertions.assertEquals(documents.size(), 1);

}

因此我们可以使用过滤器可以实现权限的控制等等。

总结

本篇文章稍微介绍了向量、Spring AI框架支持的向量数据库,对PgVector数据库进行了安装,并使用。最后使用Spring AI 框架实现测试代码完成插入、查询等演示,总体上比较简单。

技术博客阅读列表: