test/agents-flex
1
0
Fork 0
mirror of https://gitee.com/agents-flex/agents-flex.git synced 2024-12-02 03:48:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: refactor DocumentStore.java

Browse Source
This commit is contained in:
Michael Yang 2024-04-25 18:36:39 +08:00
parent da71211374
commit 5b2eaa64fb
6 changed files with 129 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

76
agents-flex-core/src/main/java/com/agentsflex/store/DocumentStore.java
View File

@ -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<Document>{
public abstract class DocumentStore extends VectorStore<Document> {
/**
* embeddings 模型，可以使用外部的 embeddings 模型，也可以使用自己的 embeddings
@ -48,5 +51,76 @@ public abstract class DocumentStore extends VectorStore<Document>{
this.documentSplitter = documentSplitter;
}
@Override
public StoreResult store(List<Document> 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<String> ids, StoreOptions options) {
if (options == null) {
options = StoreOptions.DEFAULT;
}
return deleteInternal(ids, options);
}
@Override
public StoreResult update(List<Document> documents, StoreOptions options) {
if (options == null) {
options = StoreOptions.DEFAULT;
}
embedDocumentsIfNecessary(documents, options);
return updateInternal(documents, options);
}
@Override
public List<Document> 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<Document> 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<Document> documents, StoreOptions options);
public abstract StoreResult deleteInternal(Collection<String> ids, StoreOptions options);
public abstract StoreResult updateInternal(List<Document> documents, StoreOptions options);
public abstract List<Document> searchInternal(SearchWrapper wrapper, StoreOptions options);
}

6
agents-flex-core/src/main/java/com/agentsflex/store/SearchWrapper.java
View File

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

27
agents-flex-core/src/main/java/com/agentsflex/store/StoreOptions.java
View File

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

18
agents-flex-core/src/main/java/com/agentsflex/store/VectorStore.java
View File

@ -26,22 +26,38 @@ import java.util.List;
*/
public abstract class VectorStore<T extends VectorData> {
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<T> documents) {
return store(documents, StoreOptions.DEFAULT);
}
public abstract StoreResult store(List<T> documents, StoreOptions options);
public StoreResult delete(Collection<String> ids) {
return delete(ids, null);
return delete(ids, StoreOptions.DEFAULT);
}
public abstract StoreResult delete(Collection<String> 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<T> documents) {
return update(documents, StoreOptions.DEFAULT);
}
public abstract StoreResult update(List<T> documents, StoreOptions options);
public abstract List<T> search(SearchWrapper wrapper, StoreOptions options);

12
agents-flex-store/agents-flex-store-aliyun/src/main/java/com/agentsflex/store/aliyun/AliyunVectorStore.java
View File

@ -43,7 +43,7 @@ public class AliyunVectorStore extends DocumentStore {
}
@Override
public StoreResult store(List<Document> documents, StoreOptions options) {
public StoreResult storeInternal(List<Document> 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<String> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<String> ids, StoreOptions options) {
Map<String, String> headers = new HashMap<>();
headers.put("Content-Type", "application/json");
@ -92,7 +92,7 @@ public class AliyunVectorStore extends DocumentStore {
@Override
public StoreResult update(List<Document> documents, StoreOptions options) {
public StoreResult updateInternal(List<Document> documents, StoreOptions options) {
if (documents == null || documents.isEmpty()) {
return StoreResult.DEFAULT_SUCCESS;
}
@ -124,7 +124,7 @@ public class AliyunVectorStore extends DocumentStore {
@Override
public List<Document> search(SearchWrapper wrapper, StoreOptions options) {
public List<Document> searchInternal(SearchWrapper wrapper, StoreOptions options) {
Map<String, String> 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<String, Object> 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;
}

8
agents-flex-store/agents-flex-store-qcloud/src/main/java/com/agentsflex/store/qcloud/QCloudVectorStore.java
View File

@ -45,7 +45,7 @@ public class QCloudVectorStore extends DocumentStore {
@Override
public StoreResult store(List<Document> documents, StoreOptions options) {
public StoreResult storeInternal(List<Document> 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<String> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<String> ids, StoreOptions options) {
Map<String, String> 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<Document> documents, StoreOptions options) {
public StoreResult updateInternal(List<Document> documents, StoreOptions options) {
if (documents == null || documents.isEmpty()) {
return StoreResult.DEFAULT_SUCCESS;
}
@ -126,7 +126,7 @@ public class QCloudVectorStore extends DocumentStore {
}
@Override
public List<Document> search(SearchWrapper searchWrapper, StoreOptions options) {
public List<Document> searchInternal(SearchWrapper searchWrapper, StoreOptions options) {
Map<String, String> headers = new HashMap<>();
headers.put("Content-Type", "application/json");
headers.put("Authorization", "Bearer account=" + config.getAccount() + "&api_key=" + config.getApiKey());