Spring Boot框架使用chromadb向量库

引入maven依赖

<!-- 这个是向量库的包 -->
<dependency>
    <groupId>io.github.amikos-tech</groupId>
    <artifactId>chromadb-java-client</artifactId>
    <version>0.1.7</version>
</dependency>

<!-- 这个是阿里的包 -->
<dependency>
    <groupId>com.alibaba</groupId>
    <artifactId>dashscope-sdk-java</artifactId>
    <version>2.16.10</version>
</dependency>

默认嵌入函数

下面是人家写好的例子，可以直接用，我复制过来的，文章最后是源码地址和文档地址

package tech.amikos;

import tech.amikos.chromadb.*;
import tech.amikos.chromadb.Collection;
import tech.amikos.chromadb.embeddings.DefaultEmbeddingFunction;

import java.util.*;

public class Main {
    public static void main(String[] args) {
        try {
            Client client = new Client(System.getenv("CHROMA_URL"));
            client.reset();
            EmbeddingFunction ef = new DefaultEmbeddingFunction();
            Collection collection = client.createCollection("test-collection", null, true, ef);
            List<Map<String, String>> metadata = new ArrayList<>();
            metadata.add(new HashMap<String, String>() {{
                put("type", "scientist");
            }});
            metadata.add(new HashMap<String, String>() {{
                put("type", "spy");
            }});
            collection.add(null, metadata, Arrays.asList("Hello, my name is John. I am a Data Scientist.", "Hello, my name is Bond. I am a Spy."), Arrays.asList("1", "2"));
            Collection.QueryResponse qr = collection.query(Arrays.asList("Who is the spy"), 10, null, null, null);
            System.out.println(qr);
        } catch (Exception e) {
            System.out.println(e);
        }
    }
}

重写嵌入函数，使用自定义的嵌入模型，这里我用的是阿里的文本向量模型

由于chromadb是国外的，默认的嵌入函数对中文不是那么友好，向量中文的内容，查询出来的结果不是那么理想，所以这里我使用阿里的文本向量模型，对中文进行向量转换，然后使用chromadb的向量模型进行存储，查询。当然，文章最后的链接文档中也有使用别的向量模型，比如“OpenAI、Cohere ”等，大家可以根据自己的需求进行选择。我这里只是阿里的文本向量模型例子，不多说，上代码：

• 1.定义了一个名为 LingJiEmbeddingFunction 的类，实现 EmbeddingFunction 接口，

public class LingJiEmbeddingFunction implements EmbeddingFunction {
    private static String key0 = "你自己的阿里key";
    private static String model = "text-embedding-v1"; // 阿里模型名称

    private TextEmbedding textEmbedding;

• 2.实现下面 EmbeddingFunction接口的方法

    Embedding embedQuery(String query) throws EFException;

    List<Embedding> embedDocuments(List<String> documents) throws EFException;

    List<Embedding> embedDocuments(String[] documents) throws EFException;

算了，第一次写文章，直接上代码吧，这里就不一一解释了，代码中有注释，大家看代码吧，有问题可以留言，我尽量回复：

    public LingJiEmbeddingFunction() {
        // 初始化灵积的TextEmbedding客户端
        this.textEmbedding = new TextEmbedding();
    }

    @Override
    public Embedding embedQuery(String query) throws EFException {
        try {
            TextEmbeddingParam param = TextEmbeddingParam
                    .builder()
                    .apiKey(key0)
                    .model(model)
                    .texts(Arrays.asList(query))
                    .build();
            TextEmbeddingResult result = textEmbedding.call(param);
            return extractEmbedding(result.getOutput().getEmbeddings().get(0));
        } catch (ApiException | NoApiKeyException e) {
            throw new EFException(e);
        }
    }

    @Override
    public List<Embedding> embedDocuments(List<String> documents) throws EFException {
        try {
            TextEmbeddingParam param = TextEmbeddingParam
                    .builder()
                    .apiKey(key0)
                    .model(model)
                    .texts(documents)
                    .build();
            TextEmbeddingResult result = textEmbedding.call(param);
            return result.getOutput().getEmbeddings().stream()
                    .map(this::extractEmbedding)
                    .collect(Collectors.toList());
        } catch (ApiException | NoApiKeyException e) {
            throw new EFException(e);
        }
    }

    @Override
    public List<Embedding> embedDocuments(String[] documents) throws EFException {
        // 直接委托给List<String>版本的实现
        return embedDocuments(Arrays.asList(documents));
    }

    private Embedding extractEmbedding(TextEmbeddingResultItem embeddingItem) {
        List<Double> vector = embeddingItem.getEmbedding();
        return new Embedding(vector);
    }

上面就是自定义使用阿里的文本向量模型了，下面是怎么使用写好的自定义模型去向量文本入库，我这只是简单使用，复杂的自己点开看sdk的接口可以实现什么：

@Slf4j
public class ChromaDBDiseaseUtils {

    private static final String CHROMA_URL = "http://127.0.0.1:31282"; // 你安装的ChromaDB地址，我这是安装在我自己电脑
    private static final String collectionName = "collectionNameTest"; // 自己定义集合名称，相对于mysql中的表名称，我是这样理解的
    private static Client client;
    private static Collection collection;

    static {
        try {
            // 创建 ChromaDB 客户端
            client = new Client(CHROMA_URL);
            // 创建自定义嵌入函数实例
            LingJiEmbeddingFunction customEF = new LingJiEmbeddingFunction();
            // 创建集合
            collection = client.createCollection(collectionName, null, true, customEF);
        } catch (Exception e) {
            log.error("初始化ChromaDB客户端或集合失败: {}", e.getMessage());
        }
    }

    /**
     * 向向量库中添加文档
     * @param documents
     * @param metadata
     * @param ids
     */
    public static void addDocuments(List<String> documents, List<Map<String, String>> metadata, List<String> ids) {
        try {
            // 向集合中添加文档
            collection.add(null, metadata, documents, ids);
        } catch (Exception e) {
            log.error("ChromaDB库添加文档失败: {}", e.getMessage());
        }
    }

    /**
     * 根据查询条件查询文档
     * @param query 查询语句
     * @param nResults 返回结果数量
     * @param where 查询条件（元数据的查询条件）
     * @return
     */
    public static Collection.QueryResponse queryDocuments(String query, Integer nResults, Map<String, Object> where) {
        try {
            // 查询文档
            return collection.query(Arrays.asList(query), nResults, where, null, null);
        } catch (Exception e) {
            log.error("ChromaDB库查询单据失败: {}", e.getMessage());
            return null;
        }
    }

    /**
     * 删除集合
     * @return
     */
    public static Collection deleteCollection(){
        try {
            // 删除集合
            return client.deleteCollection(collectionName);
        } catch (Exception e) {
            log.error("ChromaDB库删除集合失败: {}", e.getMessage());
            return null;
        }
    }

}

工具写好了，那么下面就是怎么使用了：

 public static void main(String[] args) {
        // 向量文本内容
        List<String> documents = new ArrayList<>();
        // 元数据（这个相对于mysql的查询条件，可以用也可以不用，如果你想精确的查询到某一条就需要这个了）
        List<Map<String, String>> metadata = new ArrayList<>();
        // 文档id
        List<String> ids = new ArrayList<>();
        documents.add("测试文本1");
        Map<String, String> meta = new HashMap<>();
        meta.put("id", "id1");
        metadata.add(meta);
        ids.add("1");
        ChromaDBDiseaseUtils.addDocuments(documents, metadata, ids);
        
        // 查询
        Collection.QueryResponse qr = ChromaDBDiseaseUtils.queryDocuments("文本", 3,null);
        // 我是查询3条，这边是打印查询出来的接口，
        qr.getDocuments().get(0).stream().forEach(System.out::println);
        qr.getMetadatas().get(0).stream().forEach(System.out::println);
        qr.getDistances().get(0).stream().forEach(System.out::println);

        // 删除集合
        // ChromaDBDiseaseUtils.deleteCollection();
    }

这样就可以使用chromadb了，当然，这只是简单的使用，如果需要更复杂的操作，可以参考chromadb的官方文档。

参考文献：GitHub - amikos-tech/chromadb-java-client: A thin client for Chroma Vector DB implemented in Java
感谢大佬们的无私奉献，才有我们这些小白可以快速上手。