test/agents-flex
1
0
Fork 0
mirror of https://gitee.com/agents-flex/agents-flex.git synced 2024-11-29 18:38:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor: optimize VectorStore

Browse Source
This commit is contained in:
Michael Yang 2024-10-26 13:04:03 +08:00
parent 3a403c2556
commit 6bb3144942
8 changed files with 40 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -88,7 +88,7 @@ public abstract class DocumentStore extends VectorStore<Document> {
}
@Override
public StoreResult delete(Collection<Object> ids, StoreOptions options) {
public StoreResult delete(Collection<?> ids, StoreOptions options) {
if (options == null) {
options = StoreOptions.DEFAULT;
}
@ -140,7 +140,7 @@ public abstract class DocumentStore extends VectorStore<Document> {
public abstract StoreResult storeInternal(List<Document> documents, StoreOptions options);
public abstract StoreResult deleteInternal(Collection<Object> ids, StoreOptions options);
public abstract StoreResult deleteInternal(Collection<?> ids, StoreOptions options);
public abstract StoreResult updateInternal(List<Document> documents, StoreOptions options);

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

@ -75,6 +75,20 @@ public abstract class VectorStore<T extends VectorData> {
* @return store result
*/
public StoreResult delete(Object... ids) {
if (ids.length == 0) {
throw new IllegalArgumentException("ids can not be empty.");
}
if (ids.length == 1) {
Object object = ids[0];
if (object == null) {
throw new NullPointerException("id can not be null.");
}
if (object instanceof Collection) {
return delete((Collection<?>) object, StoreOptions.DEFAULT);
}
}
return delete(Arrays.asList(ids), StoreOptions.DEFAULT);
}
@ -85,7 +99,7 @@ public abstract class VectorStore<T extends VectorData> {
* @param ids the ids
* @return store result
*/
public StoreResult delete(Collection<Object> ids) {
public StoreResult delete(Collection<?> ids) {
return delete(ids, StoreOptions.DEFAULT);
}
@ -96,7 +110,7 @@ public abstract class VectorStore<T extends VectorData> {
* @param options store options
* @return store result
*/
public abstract StoreResult delete(Collection<Object> ids, StoreOptions options);
public abstract StoreResult delete(Collection<?> ids, StoreOptions options);
/**
* update the vector data by id

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

@ -89,7 +89,7 @@ public class AliyunVectorStore extends DocumentStore {
@Override
public StoreResult deleteInternal(Collection<Object> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
Map<String, String> headers = new HashMap<>();
headers.put("Content-Type", "application/json");
headers.put("dashvector-auth-token", config.getApiKey());

2
agents-flex-store/agents-flex-store-elasticsearch/src/main/java/com/agentsflex/store/elasticsearch/ElasticSearchVectorStore.java
View File

@ -141,7 +141,7 @@ public class ElasticSearchVectorStore extends DocumentStore {
}
@Override
public StoreResult deleteInternal(Collection<Object> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
BulkRequest.Builder bulkBuilder = new BulkRequest.Builder();
for (Object id : ids) {

2
agents-flex-store/agents-flex-store-milvus/src/main/java/com/agentsflex/store/milvus/MilvusVectorStore.java
View File

@ -185,7 +185,7 @@ public class MilvusVectorStore extends DocumentStore {
}
@Override
public StoreResult deleteInternal(Collection<Object> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
DeleteReq.DeleteReqBuilder<?, ?> builder = DeleteReq.builder();
if (StringUtil.hasText(options.getPartitionName())) {

2
agents-flex-store/agents-flex-store-opensearch/src/main/java/com/agentsflex/store/opensearch/OpenSearchVectorStore.java
View File

@ -201,7 +201,7 @@ public class OpenSearchVectorStore extends DocumentStore {
}
@Override
public StoreResult deleteInternal(Collection<Object> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
String indexName = options.getIndexNameOrDefault(config.getDefaultIndexName());
BulkRequest.Builder bulkBuilder = new BulkRequest.Builder();
for (Object id : ids) {

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

@ -78,7 +78,7 @@ public class QCloudVectorStore extends DocumentStore {
@Override
public StoreResult deleteInternal(Collection<Object> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<?> 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());

26
agents-flex-store/agents-flex-store-redis/src/main/java/com/agentsflex/store/redis/RedisVectorStore.java
View File

@ -23,9 +23,15 @@ import com.agentsflex.core.store.StoreResult;
import com.agentsflex.core.util.StringUtil;
import com.alibaba.fastjson.JSON;
import kotlin.collections.ArrayDeque;
import redis.clients.jedis.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import redis.clients.jedis.JedisPooled;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.json.Path2;
import redis.clients.jedis.search.*;
import redis.clients.jedis.search.FTCreateParams;
import redis.clients.jedis.search.IndexDataType;
import redis.clients.jedis.search.Query;
import redis.clients.jedis.search.SearchResult;
import redis.clients.jedis.search.schemafields.SchemaField;
import redis.clients.jedis.search.schemafields.TextField;
import redis.clients.jedis.search.schemafields.VectorField;
@ -34,7 +40,6 @@ import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.text.MessageFormat;
import java.util.*;
public class RedisVectorStore extends DocumentStore {
@ -42,6 +47,8 @@ public class RedisVectorStore extends DocumentStore {
private final RedisVectorStoreConfig config;
private final JedisPooled jedis;
private final Set<String> redisIndexesCache = new HashSet<>();
private static final Logger logger = LoggerFactory.getLogger(RedisVectorStore.class);
public RedisVectorStore(RedisVectorStoreConfig config) {
this.config = config;
@ -120,8 +127,8 @@ public class RedisVectorStore extends DocumentStore {
List<Object> objects = pipeline.syncAndReturnAll();
for (Object object : objects) {
if (!object.equals("OK")) {
String message = MessageFormat.format("Could not store document: {0}", object);
throw new RuntimeException(message);
logger.error("Could not store document: {}", object);
return StoreResult.fail();
}
}
}
@ -131,7 +138,7 @@ public class RedisVectorStore extends DocumentStore {
@Override
public StoreResult deleteInternal(Collection<Object> ids, StoreOptions options) {
public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
try (Pipeline pipeline = this.jedis.pipelined()) {
for (Object id : ids) {
String key = config.getStorePrefix() + id;
@ -141,8 +148,8 @@ public class RedisVectorStore extends DocumentStore {
List<Object> objects = pipeline.syncAndReturnAll();
for (Object object : objects) {
if (!object.equals(1L)) {
String message = MessageFormat.format("Could not delete document: {0}", object);
throw new RuntimeException(message);
logger.error("Could not delete document: {}", object);
return StoreResult.fail();
}
}
}
@ -178,7 +185,8 @@ public class RedisVectorStore extends DocumentStore {
Query query = new Query("*=>[KNN " + wrapper.getMaxResults() + " @vector $BLOB AS score]")
.addParam("BLOB", vectorBytes)
.returnFields("text", "vector", "score")
.setSortBy("score", true)
// 按照 score 排序，降序, 相似度越高的越靠前
.setSortBy("score", false)
.dialect(2);
// 执行搜索