Add File
This commit is contained in:
@@ -0,0 +1,158 @@
|
||||
/*
|
||||
* Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com).
|
||||
* <p>
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
* <p>
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
* <p>
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package com.agentsflex.store.vectorex;
|
||||
|
||||
import com.agentsflex.core.document.Document;
|
||||
import com.agentsflex.core.store.DocumentStore;
|
||||
import com.agentsflex.core.store.SearchWrapper;
|
||||
import com.agentsflex.core.store.StoreOptions;
|
||||
import com.agentsflex.core.store.StoreResult;
|
||||
import com.agentsflex.core.util.VectorUtil;
|
||||
import io.github.javpower.vectorexclient.VectorRexClient;
|
||||
import io.github.javpower.vectorexclient.builder.QueryBuilder;
|
||||
import io.github.javpower.vectorexclient.entity.MetricType;
|
||||
import io.github.javpower.vectorexclient.entity.ScalarField;
|
||||
import io.github.javpower.vectorexclient.entity.VectoRexEntity;
|
||||
import io.github.javpower.vectorexclient.entity.VectorFiled;
|
||||
import io.github.javpower.vectorexclient.req.CollectionDataAddReq;
|
||||
import io.github.javpower.vectorexclient.req.CollectionDataDelReq;
|
||||
import io.github.javpower.vectorexclient.req.VectoRexCollectionReq;
|
||||
import io.github.javpower.vectorexclient.res.DbData;
|
||||
import io.github.javpower.vectorexclient.res.ServerResponse;
|
||||
import io.github.javpower.vectorexclient.res.VectorSearchResult;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import java.util.*;
|
||||
|
||||
public class VectoRexStore extends DocumentStore {
|
||||
|
||||
private final VectoRexStoreConfig config;
|
||||
private final VectorRexClient client;
|
||||
private final String defaultCollectionName;
|
||||
private boolean isCreateCollection=false;
|
||||
private static final Logger logger = LoggerFactory.getLogger(VectoRexStore.class);
|
||||
|
||||
public VectoRexStore(VectoRexStoreConfig config) {
|
||||
this.config = config;
|
||||
this.defaultCollectionName=config.getDefaultCollectionName();
|
||||
this.client = new VectorRexClient(config.getUri(),config.getUsername(),config.getPassword());
|
||||
}
|
||||
|
||||
@Override
|
||||
public StoreResult storeInternal(List<Document> documents, StoreOptions options) {
|
||||
List<Map<String, Object>> data=new ArrayList<>();
|
||||
for (Document doc : documents) {
|
||||
Map<String, Object> dict=new HashMap<>();
|
||||
dict.put("id",String.valueOf(doc.getId()));
|
||||
dict.put("content", doc.getContent());
|
||||
dict.put("vector", VectorUtil.toFloatList(doc.getVector()));
|
||||
data.add(dict);
|
||||
}
|
||||
String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
|
||||
if(config.isAutoCreateCollection()&&!isCreateCollection){
|
||||
ServerResponse<List<VectoRexEntity>> collections = client.getCollections();
|
||||
if (collections.getData() == null || collections.getData().stream().noneMatch(e -> e.getCollectionName().equals(collectionName))) {
|
||||
createCollection(collectionName);
|
||||
}else {
|
||||
isCreateCollection=true;
|
||||
}
|
||||
}
|
||||
for (Map<String, Object> map : data) {
|
||||
CollectionDataAddReq req = CollectionDataAddReq.builder().collectionName(collectionName).metadata(map).build();
|
||||
client.addCollectionData(req);
|
||||
}
|
||||
return StoreResult.successWithIds(documents);
|
||||
}
|
||||
|
||||
|
||||
private Boolean createCollection(String collectionName) {
|
||||
List<ScalarField> scalarFields = new ArrayList();
|
||||
ScalarField id = ScalarField.builder().name("id").isPrimaryKey(true).build();
|
||||
ScalarField content = ScalarField.builder().name("content").isPrimaryKey(false).build();
|
||||
scalarFields.add(id);
|
||||
scalarFields.add(content);
|
||||
List<VectorFiled> vectorFiles = new ArrayList();
|
||||
VectorFiled vector = VectorFiled.builder().name("vector").metricType(MetricType.FLOAT_COSINE_DISTANCE).dimensions(this.getEmbeddingModel().dimensions()).build();
|
||||
vectorFiles.add(vector);
|
||||
ServerResponse<Void> response = client.createCollection(VectoRexCollectionReq.builder().collectionName(collectionName).scalarFields(scalarFields).vectorFileds(vectorFiles).build());
|
||||
return response.isSuccess();
|
||||
}
|
||||
|
||||
@Override
|
||||
public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
|
||||
for (Object id : ids) {
|
||||
CollectionDataDelReq req = CollectionDataDelReq.builder().collectionName(options.getCollectionNameOrDefault(defaultCollectionName)).id((String) id).build();
|
||||
ServerResponse<Void> response = client.deleteCollectionData(req);
|
||||
if (!response.isSuccess()) {
|
||||
return StoreResult.fail();
|
||||
}
|
||||
}
|
||||
return StoreResult.success();
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> searchInternal(SearchWrapper searchWrapper, StoreOptions options) {
|
||||
ServerResponse<List<VectorSearchResult>> response = client.queryCollectionData(QueryBuilder.lambda(options.getCollectionNameOrDefault(defaultCollectionName)).vector("vector", Collections.singletonList(VectorUtil.toFloatList(searchWrapper.getVector()))).topK(searchWrapper.getMaxResults()));
|
||||
if(!response.isSuccess()){
|
||||
logger.error("Error searching in VectoRex", response.getMsg());
|
||||
return Collections.emptyList();
|
||||
}
|
||||
List<VectorSearchResult> data = response.getData();
|
||||
List<Document> documents = new ArrayList<>();
|
||||
for (VectorSearchResult result : data) {
|
||||
DbData dd = result.getData();
|
||||
Map<String, Object> metadata = dd.getMetadata();
|
||||
Document doc=new Document();
|
||||
doc.setId(result.getId());
|
||||
doc.setContent((String) metadata.get("content"));
|
||||
Object vectorObj = metadata.get("vector");
|
||||
if (vectorObj instanceof List) {
|
||||
//noinspection unchecked
|
||||
doc.setVector(VectorUtil.convertToVector((List<Float>) vectorObj));
|
||||
}
|
||||
documents.add(doc);
|
||||
}
|
||||
return documents;
|
||||
}
|
||||
|
||||
@Override
|
||||
public StoreResult updateInternal(List<Document> documents, StoreOptions options) {
|
||||
if (documents == null || documents.isEmpty()) {
|
||||
return StoreResult.success();
|
||||
}
|
||||
List<Map<String, Object>> data=new ArrayList<>();
|
||||
for (Document doc : documents) {
|
||||
Map<String, Object> dict=new HashMap<>();
|
||||
dict.put("id",String.valueOf(doc.getId()));
|
||||
dict.put("content", doc.getContent());
|
||||
dict.put("vector", VectorUtil.toFloatList(doc.getVector()));
|
||||
data.add(dict);
|
||||
}
|
||||
String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
|
||||
for (Map<String, Object> map : data) {
|
||||
CollectionDataAddReq req = CollectionDataAddReq.builder().collectionName(collectionName).metadata(map).build();
|
||||
client.updateCollectionData(req);
|
||||
}
|
||||
return StoreResult.successWithIds(documents);
|
||||
}
|
||||
|
||||
public VectorRexClient getClient() {
|
||||
return client;
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user