From 5b2eaa64fb6416f054498772fa6f2d4e913b76a9 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Thu, 25 Apr 2024 18:36:39 +0800 Subject: [PATCH] refactor: refactor DocumentStore.java --- .../com/agentsflex/store/DocumentStore.java | 76 ++++++++++++++++++- .../com/agentsflex/store/SearchWrapper.java | 6 +- .../com/agentsflex/store/StoreOptions.java | 27 ++++++- .../com/agentsflex/store/VectorStore.java | 18 ++++- .../store/aliyun/AliyunVectorStore.java | 12 +-- .../store/qcloud/QCloudVectorStore.java | 8 +- 6 files changed, 129 insertions(+), 18 deletions(-) diff --git a/agents-flex-core/src/main/java/com/agentsflex/store/DocumentStore.java b/agents-flex-core/src/main/java/com/agentsflex/store/DocumentStore.java index 6f14618..a2f5d72 100644 --- a/agents-flex-core/src/main/java/com/agentsflex/store/DocumentStore.java +++ b/agents-flex-core/src/main/java/com/agentsflex/store/DocumentStore.java @@ -19,10 +19,13 @@ import com.agentsflex.document.Document; import com.agentsflex.document.DocumentSplitter; import com.agentsflex.llm.embedding.EmbeddingModel; +import java.util.Collection; +import java.util.List; + /** * 文档存储器 */ -public abstract class DocumentStore extends VectorStore{ +public abstract class DocumentStore extends VectorStore { /** * embeddings 模型，可以使用外部的 embeddings 模型，也可以使用自己的 embeddings @@ -48,5 +51,76 @@ public abstract class DocumentStore extends VectorStore{ this.documentSplitter = documentSplitter; } + @Override + public StoreResult store(List documents, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + if (documentSplitter != null) { + documents = documentSplitter.splitAll(documents); + } + + embedDocumentsIfNecessary(documents, options); + + return storeInternal(documents, options); + } + + @Override + public StoreResult delete(Collection ids, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + return deleteInternal(ids, options); + } + + @Override + public StoreResult update(List documents, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + + embedDocumentsIfNecessary(documents, options); + return updateInternal(documents, options); + } + + + @Override + public List search(SearchWrapper wrapper, StoreOptions options) { + if (options == null) { + options = StoreOptions.DEFAULT; + } + + if (wrapper.getVector() == null && embeddingModel != null && wrapper.isWithVector()) { + VectorData vectorData = embeddingModel.embed(Document.of(wrapper.getText()), options.getEmbeddingOptions()); + if (vectorData != null) { + wrapper.setVector(vectorData.getVector()); + } + } + + return searchInternal(wrapper, options); + } + + + protected void embedDocumentsIfNecessary(List documents, StoreOptions options) { + if (embeddingModel != null) { + for (Document document : documents) { + if (document.getVector() == null) { + VectorData vectorData = embeddingModel.embed(document, options.getEmbeddingOptions()); + if (vectorData != null) { + document.setVector(vectorData.getVector()); + } + } + } + } + } + + + public abstract StoreResult storeInternal(List documents, StoreOptions options); + + public abstract StoreResult deleteInternal(Collection ids, StoreOptions options); + + public abstract StoreResult updateInternal(List documents, StoreOptions options); + + public abstract List searchInternal(SearchWrapper wrapper, StoreOptions options); } diff --git a/agents-flex-core/src/main/java/com/agentsflex/store/SearchWrapper.java b/agents-flex-core/src/main/java/com/agentsflex/store/SearchWrapper.java index 96db34d..2a3b954 100644 --- a/agents-flex-core/src/main/java/com/agentsflex/store/SearchWrapper.java +++ b/agents-flex-core/src/main/java/com/agentsflex/store/SearchWrapper.java @@ -46,7 +46,7 @@ public class SearchWrapper extends VectorData { /** * 是否包含向量数据查询，如果当前值为 true，且向量内容为 null 时，会自动通过向量数据库把 text 转换为 向量数据 */ - private Boolean withVector; + private boolean withVector = true; /** * 查询条件 @@ -103,11 +103,11 @@ public class SearchWrapper extends VectorData { return this; } - public Boolean getWithVector() { + public boolean isWithVector() { return withVector; } - public void setWithVector(Boolean withVector) { + public void setWithVector(boolean withVector) { this.withVector = withVector; } diff --git a/agents-flex-core/src/main/java/com/agentsflex/store/StoreOptions.java b/agents-flex-core/src/main/java/com/agentsflex/store/StoreOptions.java index 7ef64fc..349829d 100644 --- a/agents-flex-core/src/main/java/com/agentsflex/store/StoreOptions.java +++ b/agents-flex-core/src/main/java/com/agentsflex/store/StoreOptions.java @@ -1,18 +1,19 @@ package com.agentsflex.store; +import com.agentsflex.llm.embedding.EmbeddingOptions; import com.agentsflex.util.StringUtil; public class StoreOptions { - public static StoreOptions EMPTY = new StoreOptions(){ + public static final StoreOptions DEFAULT = new StoreOptions() { @Override public void setCollectionName(String collectionName) { - throw new IllegalStateException("Can not set collectionName to the empty instance."); + throw new IllegalStateException("Can not set collectionName to the default instance."); } @Override public void setPartitionName(String partitionName) { - throw new IllegalStateException("Can not set partitionName to the empty instance."); + throw new IllegalStateException("Can not set partitionName to the default instance."); } }; @@ -26,6 +27,11 @@ public class StoreOptions { */ private String partitionName; + /** + * embedding 的配置内容 + */ + private EmbeddingOptions embeddingOptions; + public String getCollectionName() { return collectionName; @@ -50,4 +56,19 @@ public class StoreOptions { public void setPartitionName(String partitionName) { this.partitionName = partitionName; } + + public EmbeddingOptions getEmbeddingOptions() { + return embeddingOptions; + } + + public void setEmbeddingOptions(EmbeddingOptions embeddingOptions) { + this.embeddingOptions = embeddingOptions; + } + + + public static StoreOptions ofCollectionName(String collectionName) { + StoreOptions storeOptions = new StoreOptions(); + storeOptions.setCollectionName(collectionName); + return storeOptions; + } } diff --git a/agents-flex-core/src/main/java/com/agentsflex/store/VectorStore.java b/agents-flex-core/src/main/java/com/agentsflex/store/VectorStore.java index f950ab7..afe0716 100644 --- a/agents-flex-core/src/main/java/com/agentsflex/store/VectorStore.java +++ b/agents-flex-core/src/main/java/com/agentsflex/store/VectorStore.java @@ -26,22 +26,38 @@ import java.util.List; */ public abstract class VectorStore { + public StoreResult store(T document) { + return store(document, StoreOptions.DEFAULT); + } + public StoreResult store(T document, StoreOptions options) { return store(Collections.singletonList(document), options); } + public StoreResult store(List documents) { + return store(documents, StoreOptions.DEFAULT); + } + public abstract StoreResult store(List documents, StoreOptions options); public StoreResult delete(Collection ids) { - return delete(ids, null); + return delete(ids, StoreOptions.DEFAULT); } public abstract StoreResult delete(Collection ids, StoreOptions options); + public StoreResult update(T document) { + return update(document, StoreOptions.DEFAULT); + } + public StoreResult update(T document, StoreOptions options) { return update(Collections.singletonList(document), options); } + public StoreResult update(List documents) { + return update(documents, StoreOptions.DEFAULT); + } + public abstract StoreResult update(List documents, StoreOptions options); public abstract List search(SearchWrapper wrapper, StoreOptions options); diff --git a/agents-flex-store/agents-flex-store-aliyun/src/main/java/com/agentsflex/store/aliyun/AliyunVectorStore.java b/agents-flex-store/agents-flex-store-aliyun/src/main/java/com/agentsflex/store/aliyun/AliyunVectorStore.java index 5fa6591..55289b9 100644 --- a/agents-flex-store/agents-flex-store-aliyun/src/main/java/com/agentsflex/store/aliyun/AliyunVectorStore.java +++ b/agents-flex-store/agents-flex-store-aliyun/src/main/java/com/agentsflex/store/aliyun/AliyunVectorStore.java @@ -43,7 +43,7 @@ public class AliyunVectorStore extends DocumentStore { } @Override - public StoreResult store(List documents, StoreOptions options) { + public StoreResult storeInternal(List documents, StoreOptions options) { if (documents == null || documents.isEmpty()) { return StoreResult.DEFAULT_SUCCESS; } @@ -75,7 +75,7 @@ public class AliyunVectorStore extends DocumentStore { @Override - public StoreResult delete(Collection ids, StoreOptions options) { + public StoreResult deleteInternal(Collection ids, StoreOptions options) { Map headers = new HashMap<>(); headers.put("Content-Type", "application/json"); @@ -92,7 +92,7 @@ public class AliyunVectorStore extends DocumentStore { @Override - public StoreResult update(List documents, StoreOptions options) { + public StoreResult updateInternal(List documents, StoreOptions options) { if (documents == null || documents.isEmpty()) { return StoreResult.DEFAULT_SUCCESS; } @@ -124,7 +124,7 @@ public class AliyunVectorStore extends DocumentStore { @Override - public List search(SearchWrapper wrapper, StoreOptions options) { + public List searchInternal(SearchWrapper wrapper, StoreOptions options) { Map headers = new HashMap<>(); headers.put("Content-Type", "application/json"); headers.put("dashvector-auth-token", config.getApiKey()); @@ -132,7 +132,7 @@ public class AliyunVectorStore extends DocumentStore { Map payloadMap = new HashMap<>(); payloadMap.put("vector", wrapper.getVector()); payloadMap.put("topk", wrapper.getMaxResults()); - payloadMap.put("include_vector", wrapper.getWithVector()); + payloadMap.put("include_vector", wrapper.isWithVector()); payloadMap.put("filter", wrapper.toFilterExpression()); String payload = JSON.toJSONString(payloadMap); @@ -147,7 +147,7 @@ public class AliyunVectorStore extends DocumentStore { int code = rootObject.getIntValue("code"); if (code != 0) { //error - LoggerFactory.getLogger(AliyunVectorStore.class).error("can not search data AliyunVectorStore, code: " + code); + LoggerFactory.getLogger(AliyunVectorStore.class).error("can not search data AliyunVectorStore（code: " + code + "), message: " + rootObject.getString("message")); return null; } diff --git a/agents-flex-store/agents-flex-store-qcloud/src/main/java/com/agentsflex/store/qcloud/QCloudVectorStore.java b/agents-flex-store/agents-flex-store-qcloud/src/main/java/com/agentsflex/store/qcloud/QCloudVectorStore.java index 06a70f7..d878d21 100644 --- a/agents-flex-store/agents-flex-store-qcloud/src/main/java/com/agentsflex/store/qcloud/QCloudVectorStore.java +++ b/agents-flex-store/agents-flex-store-qcloud/src/main/java/com/agentsflex/store/qcloud/QCloudVectorStore.java @@ -45,7 +45,7 @@ public class QCloudVectorStore extends DocumentStore { @Override - public StoreResult store(List documents, StoreOptions options) { + public StoreResult storeInternal(List documents, StoreOptions options) { if (documents == null || documents.isEmpty()) { return StoreResult.DEFAULT_SUCCESS; } @@ -78,7 +78,7 @@ public class QCloudVectorStore extends DocumentStore { @Override - public StoreResult delete(Collection ids, StoreOptions options) { + public StoreResult deleteInternal(Collection ids, StoreOptions options) { Map headers = new HashMap<>(); headers.put("Content-Type", "application/json"); headers.put("Authorization", "Bearer account=" + config.getAccount() + "&api_key=" + config.getApiKey()); @@ -100,7 +100,7 @@ public class QCloudVectorStore extends DocumentStore { @Override - public StoreResult update(List documents, StoreOptions options) { + public StoreResult updateInternal(List documents, StoreOptions options) { if (documents == null || documents.isEmpty()) { return StoreResult.DEFAULT_SUCCESS; } @@ -126,7 +126,7 @@ public class QCloudVectorStore extends DocumentStore { } @Override - public List search(SearchWrapper searchWrapper, StoreOptions options) { + public List searchInternal(SearchWrapper searchWrapper, StoreOptions options) { Map headers = new HashMap<>(); headers.put("Content-Type", "application/json"); headers.put("Authorization", "Bearer account=" + config.getAccount() + "&api_key=" + config.getApiKey());