From 92bdb5ddd5f266edac884e19456e6572d545399a Mon Sep 17 00:00:00 2001 From: 0007 <0007@qq.com> Date: Wed, 27 Aug 2025 19:57:58 +0800 Subject: [PATCH] Add File --- .../store/pgvector/PgvectorDbTest.java | 134 ++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 agents-flex-store/agents-flex-store-pgvector/src/test/java/com/agentsflex/store/pgvector/PgvectorDbTest.java diff --git a/agents-flex-store/agents-flex-store-pgvector/src/test/java/com/agentsflex/store/pgvector/PgvectorDbTest.java b/agents-flex-store/agents-flex-store-pgvector/src/test/java/com/agentsflex/store/pgvector/PgvectorDbTest.java new file mode 100644 index 0000000..f2091df --- /dev/null +++ b/agents-flex-store/agents-flex-store-pgvector/src/test/java/com/agentsflex/store/pgvector/PgvectorDbTest.java @@ -0,0 +1,134 @@ +package com.agentsflex.store.pgvector; + +import com.agentsflex.core.document.Document; +import com.agentsflex.core.store.SearchWrapper; +import com.agentsflex.core.store.StoreResult; +import com.agentsflex.core.util.Maps; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +public class PgvectorDbTest { + + @Test + public void testInsert() { + PgvectorVectorStoreConfig config = new PgvectorVectorStoreConfig(); + config.setHost("127.0.0.1"); + config.setPort(5432); + config.setDatabaseName("pgvector_test"); + config.setUsername("test"); + config.setPassword("123456"); + config.setVectorDimension(1024); + config.setUseHnswIndex(true); + config.setAutoCreateCollection(true); + config.setDefaultCollectionName("test"); + + PgvectorVectorStore store = new PgvectorVectorStore(config); + Document doc = new Document("测试数据"); + // 初始化 vector 为长度为 1024 的全是 1 的数组 + double[] vector = new double[1024]; + Arrays.fill(vector, 1.0); + + doc.setVector(vector); + doc.setMetadataMap(Maps.of("test", "test")); + store.store(doc); + } + + @Test + public void testInsertMany() { + PgvectorVectorStoreConfig config = new PgvectorVectorStoreConfig(); + config.setHost("127.0.0.1"); + config.setPort(5432); + config.setDatabaseName("pgvector_test"); + config.setUsername("test"); + config.setPassword("123456"); + config.setVectorDimension(1024); + config.setUseHnswIndex(true); + config.setAutoCreateCollection(true); + config.setDefaultCollectionName("test"); + + PgvectorVectorStore store = new PgvectorVectorStore(config); + List docs = new ArrayList<>(100); + for (int i = 0; i < 100; i++) { + Document doc = new Document("测试数据" + i); + // 初始化 vector 为长度为 1024 的全是 1 的数组 + double[] vector = new double[1024]; + Arrays.fill(vector, Math.random()); + + doc.setVector(vector); + doc.setMetadataMap(Maps.of("test", "test" + i)); + docs.add(doc); + } + + store.store(docs); + } + + @Test + public void testSearch() { + PgvectorVectorStoreConfig config = new PgvectorVectorStoreConfig(); + config.setHost("127.0.0.1"); + config.setPort(5432); + config.setDatabaseName("pgvector_test"); + config.setUsername("test"); + config.setPassword("123456"); + config.setVectorDimension(1024); + config.setUseHnswIndex(true); + config.setAutoCreateCollection(true); + config.setDefaultCollectionName("test"); + PgvectorVectorStore store = new PgvectorVectorStore(config); + + double[] vector = new double[1024]; + Arrays.fill(vector, 1.0); + + SearchWrapper searchWrapper = new SearchWrapper().text("测试数据"); + searchWrapper.setVector(vector); + searchWrapper.setMinScore(0.0); + searchWrapper.setOutputVector(true); + List docs = store.search(searchWrapper); + System.out.println(docs); + } + + @Test + public void testUpdate() { + PgvectorVectorStoreConfig config = new PgvectorVectorStoreConfig(); + config.setHost("127.0.0.1"); + config.setPort(5432); + config.setDatabaseName("pgvector_test"); + config.setUsername("test"); + config.setPassword("123456"); + config.setVectorDimension(1024); + config.setUseHnswIndex(true); + config.setAutoCreateCollection(true); + config.setDefaultCollectionName("test"); + PgvectorVectorStore store = new PgvectorVectorStore(config); + + Document document = new Document("测试数据"); + document.setId("145314895749100ae8306079519b3393"); + document.setMetadataMap(Maps.of("test", "test0")); + double[] vector = new double[1024]; + Arrays.fill(vector, 1.1); + document.setVector(vector); + StoreResult update = store.update(document); + System.out.println(update); + } + + @Test + public void testDelete() { + PgvectorVectorStoreConfig config = new PgvectorVectorStoreConfig(); + config.setHost("127.0.0.1"); + config.setPort(5432); + config.setDatabaseName("pgvector_test"); + config.setUsername("test"); + config.setPassword("123456"); + config.setVectorDimension(1024); + config.setUseHnswIndex(true); + config.setAutoCreateCollection(true); + config.setDefaultCollectionName("test"); + PgvectorVectorStore store = new PgvectorVectorStore(config); + + StoreResult update = store.delete("145314895749100ae8306079519b3393","e83518d36b6d5de8199b40e3ef4e4ce1"); + System.out.println(update); + } +}