diff --git a/agents-flex-store/agents-flex-store-vectorexdb/src/main/java/com/agentsflex/store/vectorex/VectoRexStore.java b/agents-flex-store/agents-flex-store-vectorexdb/src/main/java/com/agentsflex/store/vectorex/VectoRexStore.java new file mode 100644 index 0000000..3757f54 --- /dev/null +++ b/agents-flex-store/agents-flex-store-vectorexdb/src/main/java/com/agentsflex/store/vectorex/VectoRexStore.java @@ -0,0 +1,157 @@ +/* + * Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com). + *

+ * 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 + *

+ * http://www.apache.org/licenses/LICENSE-2.0 + *

+ * 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.CollectionUtil; +import com.agentsflex.core.util.VectorUtil; +import com.google.common.collect.Lists; +import io.github.javpower.vectorex.keynote.core.DbData; +import io.github.javpower.vectorex.keynote.core.VectorData; +import io.github.javpower.vectorex.keynote.core.VectorSearchResult; +import io.github.javpower.vectorex.keynote.model.MetricType; +import io.github.javpower.vectorex.keynote.model.VectorFiled; +import io.github.javpower.vectorexcore.VectoRexClient; +import io.github.javpower.vectorexcore.entity.KeyValue; +import io.github.javpower.vectorexcore.entity.ScalarField; +import io.github.javpower.vectorexcore.entity.VectoRexEntity; + +import java.util.*; + +public class VectoRexStore extends DocumentStore { + + private final VectoRexStoreConfig config; + private final VectoRexClient client; + private final String defaultCollectionName; + private boolean isCreateCollection=false; + + public VectoRexStore(VectoRexStoreConfig config) { + this.config = config; + this.defaultCollectionName=config.getDefaultCollectionName(); + this.client = new VectoRexClient(config.getUri()); + } + @Override + public StoreResult storeInternal(List documents, StoreOptions options) { + List data=new ArrayList<>(); + for (Document doc : documents) { + Map dict=new HashMap<>(); + dict.put("id",String.valueOf(doc.getId())); + dict.put("content", doc.getContent()); + dict.put("vector", VectorUtil.toFloatList(doc.getVector())); + DbData dbData=new DbData(); + dbData.setId(String.valueOf(doc.getId())); + dbData.setMetadata(dict); + VectorData vd=new VectorData(dbData.getId(),VectorUtil.toFloatArray(doc.getVector())); + vd.setName("vector"); + dbData.setVectorFiled(Lists.newArrayList(vd)); + data.add(dbData); + } + String collectionName = options.getCollectionNameOrDefault(defaultCollectionName); + if(config.isAutoCreateCollection()&&!isCreateCollection){ + List collections = client.getCollections(); + if(CollectionUtil.noItems(collections)||collections.stream().noneMatch(e -> e.getCollectionName().equals(collectionName))){ + createCollection(collectionName); + }else { + isCreateCollection=true; + } + } + if(CollectionUtil.hasItems(data)){ + client.getStore(collectionName).saveAll(data); + } + return StoreResult.successWithIds(documents); + } + + + private void createCollection(String collectionName) { + VectorFiled vectorFiled = new VectorFiled(); + vectorFiled.setDimensions(this.getEmbeddingModel().dimensions()); + vectorFiled.setName("vector"); + vectorFiled.setMetricType(MetricType.FLOAT_COSINE_DISTANCE); + VectoRexEntity entity=new VectoRexEntity(); + entity.setCollectionName(collectionName); + List> vectorFiles=new ArrayList<>(); + vectorFiles.add(new KeyValue<>("vector",vectorFiled)); + List> scalarFields=new ArrayList<>(); + ScalarField id = new ScalarField(); + id.setName("id"); + id.setIsPrimaryKey(true); + scalarFields.add(new KeyValue<>("id",id)); + ScalarField content = new ScalarField(); + content.setName("content"); + content.setIsPrimaryKey(false); + scalarFields.add(new KeyValue<>("content",content)); + entity.setVectorFileds(vectorFiles); + entity.setScalarFields(scalarFields); + client.createCollection(entity); + } + + @Override + public StoreResult deleteInternal(Collection ids, StoreOptions options) { + client.getStore(options.getCollectionNameOrDefault(defaultCollectionName)).deleteAll((List) ids); + return StoreResult.success(); + + } + + @Override + public List searchInternal(SearchWrapper searchWrapper, StoreOptions options) { + List data = client.getStore(options.getCollectionNameOrDefault(defaultCollectionName)).search("vector", VectorUtil.toFloatList(searchWrapper.getVector()), searchWrapper.getMaxResults(), null); + List documents = new ArrayList<>(); + for (VectorSearchResult result : data) { + DbData dd = result.getData(); + Map 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) vectorObj)); + } + documents.add(doc); + } + return documents; + } + + @Override + public StoreResult updateInternal(List documents, StoreOptions options) { + if (documents == null || documents.isEmpty()) { + return StoreResult.success(); + } + for (Document doc : documents) { + Map dict=new HashMap<>(); + dict.put("id",String.valueOf(doc.getId())); + dict.put("content", doc.getContent()); + dict.put("vector", VectorUtil.toFloatList(doc.getVector())); + DbData dbData=new DbData(); + dbData.setId(String.valueOf(doc.getId())); + dbData.setMetadata(dict); + VectorData vd=new VectorData(dbData.getId(),VectorUtil.toFloatArray(doc.getVector())); + vd.setName("vector"); + dbData.setVectorFiled(Lists.newArrayList(vd)); + client.getStore(options.getCollectionNameOrDefault(defaultCollectionName)).update(dbData); + } + return StoreResult.successWithIds(documents); + } + + public VectoRexClient getClient() { + return client; + } + +}