Spring Boot框架使用chromadb向量库 - IT评测·应用市场-qidao123.com

<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);
}
}
}

复制代码

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

复制代码

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);
}

复制代码

@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();
}

复制代码