From 4d692a542be33b23c58fd14830b82c1577abc5e4 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Mon, 10 Feb 2025 17:25:06 -0500 Subject: [PATCH] Add missing integration tests for delete by ID API in vector store implementations. Extract common vector store delete tests to base class This commit extracts shared delete operation tests into a reusable BaseVectorStoreTests class. This reduces code duplication and provides a consistent test suite for delete operations across different vector store implementations. The base class includes tests for: Deleting by ID Deleting by filter expressions Deleting by string filter expressions Most of the vector store implementation now extends this base class and inherits these common tests while maintaining the ability to add vector store specific tests. Adding javadoc Signed-off-by: Soby Chacko --- .../vectorstore/BaseVectorStoreTests.java | 170 ++++++++++++++++++ .../cassandra/CassandraVectorStoreIT.java | 28 ++- .../vectorstore/ChromaVectorStoreIT.java | 75 ++------ .../ElasticsearchVectorStoreIT.java | 97 ++-------- .../vectorstore/mariadb/MariaDBStoreIT.java | 78 ++------ .../milvus/MilvusVectorStoreIT.java | 80 ++------- .../atlas/MongoDBAtlasVectorStoreIT.java | 77 ++------ .../vectorstore/neo4j/Neo4jVectorStoreIT.java | 73 ++------ .../opensearch/OpenSearchVectorStoreIT.java | 32 ++++ .../oracle/OracleVectorStoreIT.java | 80 ++------- .../vectorstore/pgvector/PgVectorStoreIT.java | 125 ++----------- .../pinecone/PineconeVectorStoreIT.java | 65 ++----- .../qdrant/QdrantVectorStoreIT.java | 63 ++----- .../vectorstore/redis/RedisVectorStoreIT.java | 65 ++----- .../typesense/TypesenseVectorStoreIT.java | 68 ++----- .../spring-ai-weaviate-store/pom.xml | 2 +- .../weaviate/WeaviateVectorStoreIT.java | 70 ++------ 17 files changed, 380 insertions(+), 868 deletions(-) create mode 100644 spring-ai-test/src/main/java/org/springframework/ai/test/vectorstore/BaseVectorStoreTests.java diff --git a/spring-ai-test/src/main/java/org/springframework/ai/test/vectorstore/BaseVectorStoreTests.java b/spring-ai-test/src/main/java/org/springframework/ai/test/vectorstore/BaseVectorStoreTests.java new file mode 100644 index 000000000..e7b7e5d23 --- /dev/null +++ b/spring-ai-test/src/main/java/org/springframework/ai/test/vectorstore/BaseVectorStoreTests.java @@ -0,0 +1,170 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * 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 + * + * https://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 org.springframework.ai.test.vectorstore; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.ai.vectorstore.filter.Filter; + +/** + * Base test class for VectorStore implementations. Provides common test scenarios for + * delete operations. + * + * @author Soby Chacko + */ +public abstract class BaseVectorStoreTests { + + /** + * Execute a test function with a configured VectorStore instance. This method is + * responsible for providing a properly initialized VectorStore within the appropriate + * Spring application context for testing. + * @param testFunction the consumer that executes test operations on the VectorStore + */ + protected abstract void executeTest(Consumer testFunction); + + protected Document createDocument(String country, Integer year) { + Map metadata = new HashMap<>(); + metadata.put("country", country); + if (year != null) { + metadata.put("year", year); + } + return new Document("The World is Big and Salvation Lurks Around the Corner", metadata); + } + + protected List setupTestDocuments(VectorStore vectorStore) { + var doc1 = createDocument("BG", 2020); + var doc2 = createDocument("NL", null); + var doc3 = createDocument("BG", 2023); + + List documents = List.of(doc1, doc2, doc3); + vectorStore.add(documents); + + return documents; + } + + private String normalizeValue(Object value) { + return value.toString().replaceAll("^\"|\"$", "").trim(); + } + + private void verifyDocumentsExist(VectorStore vectorStore, List documents) { + await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> { + List results = vectorStore.similaritySearch( + SearchRequest.builder().query("The World").topK(documents.size()).similarityThresholdAll().build()); + assertThat(results).hasSize(documents.size()); + }); + } + + private void verifyDocumentsDeleted(VectorStore vectorStore, List deletedIds) { + await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> { + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("The World").topK(10).similarityThresholdAll().build()); + + List foundIds = results.stream().map(Document::getId).collect(Collectors.toList()); + + assertThat(foundIds).doesNotContainAnyElementsOf(deletedIds); + }); + } + + @Test + protected void deleteById() { + executeTest(vectorStore -> { + List documents = setupTestDocuments(vectorStore); + verifyDocumentsExist(vectorStore, documents); + + List idsToDelete = List.of(documents.get(0).getId(), documents.get(1).getId()); + vectorStore.delete(idsToDelete); + verifyDocumentsDeleted(vectorStore, idsToDelete); + + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(documents.get(2).getId()); + Map metadata = results.get(0).getMetadata(); + assertThat(normalizeValue(metadata.get("country"))).isEqualTo("BG"); + assertThat(normalizeValue(metadata.get("year"))).isEqualTo("2023"); + + vectorStore.delete(List.of(documents.get(2).getId())); + }); + } + + @Test + protected void deleteWithStringFilterExpression() { + executeTest(vectorStore -> { + List documents = setupTestDocuments(vectorStore); + verifyDocumentsExist(vectorStore, documents); + + List bgDocIds = documents.stream() + .filter(d -> "BG".equals(d.getMetadata().get("country"))) + .map(Document::getId) + .collect(Collectors.toList()); + + vectorStore.delete("country == 'BG'"); + verifyDocumentsDeleted(vectorStore, bgDocIds); + + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); + + assertThat(results).hasSize(1); + assertThat(normalizeValue(results.get(0).getMetadata().get("country"))).isEqualTo("NL"); + + vectorStore.delete(List.of(documents.get(1).getId())); + }); + } + + @Test + protected void deleteByFilter() { + executeTest(vectorStore -> { + List documents = setupTestDocuments(vectorStore); + verifyDocumentsExist(vectorStore, documents); + + List bgDocIds = documents.stream() + .filter(d -> "BG".equals(d.getMetadata().get("country"))) + .map(Document::getId) + .collect(Collectors.toList()); + + Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, + new Filter.Key("country"), new Filter.Value("BG")); + + vectorStore.delete(filterExpression); + verifyDocumentsDeleted(vectorStore, bgDocIds); + + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); + + assertThat(results).hasSize(1); + assertThat(normalizeValue(results.get(0).getMetadata().get("country"))).isEqualTo("NL"); + + vectorStore.delete(List.of(documents.get(1).getId())); + }); + } + +} diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java index e30ba900d..7453350b2 100644 --- a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java +++ b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java @@ -19,10 +19,12 @@ package org.springframework.ai.vectorstore.cassandra; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import com.datastax.oss.driver.api.core.CqlSession; @@ -40,8 +42,10 @@ import org.springframework.ai.cassandra.CassandraImage; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.transformers.TransformersEmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn; import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumnTags; import org.springframework.ai.vectorstore.filter.Filter; @@ -64,7 +68,7 @@ import static org.assertj.core.api.Assertions.assertThat; * @since 1.0.0 */ @Testcontainers -class CassandraVectorStoreIT { +class CassandraVectorStoreIT extends BaseVectorStoreTests { @Container static CassandraContainer cassandraContainer = new CassandraContainer<>(CassandraImage.DEFAULT_IMAGE); @@ -110,6 +114,24 @@ class CassandraVectorStoreIT { return store; } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + + @Override + protected Document createDocument(String country, Integer year) { + Map metadata = new HashMap<>(); + metadata.put("country", country); + if (year != null) { + metadata.put("year", year.shortValue()); + } + return new Document("The World is Big and Salvation Lurks Around the Corner", metadata); + } + @Test void ensureBeanGetsCreated() { this.contextRunner.run(context -> { @@ -422,7 +444,7 @@ class CassandraVectorStoreIT { } @Test - void deleteByFilter() { + protected void deleteByFilter() { this.contextRunner.run(context -> { try (CassandraVectorStore store = createTestStore(context, new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED), @@ -458,7 +480,7 @@ class CassandraVectorStoreIT { } @Test - void deleteWithStringFilterExpression() { + protected void deleteWithStringFilterExpression() { this.contextRunner.run(context -> { try (CassandraVectorStore store = createTestStore(context, new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED), diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java index 765e3b563..12d4c2e00 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java @@ -20,6 +20,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.function.Consumer; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; @@ -33,6 +34,7 @@ import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -51,7 +53,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -public class ChromaVectorStoreIT { +public class ChromaVectorStoreIT extends BaseVectorStoreTests { @Container static ChromaDBContainer chromaContainer = new ChromaDBContainer(ChromaImage.DEFAULT_IMAGE); @@ -68,6 +70,14 @@ public class ChromaVectorStoreIT { "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression", Collections.singletonMap("meta2", "meta2"))); + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test public void addAndSearch() { this.contextRunner.run(context -> { @@ -168,69 +178,6 @@ public class ChromaVectorStoreIT { }); } - @Test - public void deleteWithFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - // Create test documents with different metadata - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "Bulgaria")); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "Netherlands")); - - // Add documents to the store - vectorStore.add(List.of(bgDocument, nlDocument)); - - // Verify initial state - var request = SearchRequest.builder().query("The World").topK(5).build(); - List results = vectorStore.similaritySearch(request); - assertThat(results).hasSize(2); - - // Delete document with country = Bulgaria - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("Bulgaria")); - - vectorStore.delete(filterExpression); - - // Verify Bulgaria document was deleted - results = vectorStore - .similaritySearch(SearchRequest.from(request).filterExpression("country == 'Bulgaria'").build()); - assertThat(results).isEmpty(); - - // Verify Netherlands document still exists - results = vectorStore - .similaritySearch(SearchRequest.from(request).filterExpression("country == 'Netherlands'").build()); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata().get("country")).isEqualTo("Netherlands"); - - // Clean up - vectorStore.delete(List.of(nlDocument.getId())); - }); - } - - @Test - public void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big", Map.of("country", "Bulgaria")); - var nlDocument = new Document("The World is Big", Map.of("country", "Netherlands")); - vectorStore.add(List.of(bgDocument, nlDocument)); - - var request = SearchRequest.builder().query("World").topK(5).build(); - assertThat(vectorStore.similaritySearch(request)).hasSize(2); - - vectorStore.delete("country == 'Bulgaria'"); - - var results = vectorStore.similaritySearch(request); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata().get("country")).isEqualTo("Netherlands"); - - vectorStore.delete(List.of(nlDocument.getId())); - }); - } - @Test public void documentUpdateTest() { diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java index 3acf51768..c6b3c0bb7 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java @@ -27,6 +27,7 @@ import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import co.elastic.clients.elasticsearch.ElasticsearchClient; import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord; @@ -53,7 +54,9 @@ import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter.Expression; import org.springframework.ai.vectorstore.filter.Filter.ExpressionType; import org.springframework.ai.vectorstore.filter.Filter.Key; @@ -71,7 +74,7 @@ import static org.hamcrest.Matchers.hasSize; @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -class ElasticsearchVectorStoreIT { +class ElasticsearchVectorStoreIT extends BaseVectorStoreTests { @Container private static final ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer( @@ -116,6 +119,14 @@ class ElasticsearchVectorStoreIT { }); } + @Override + protected void executeTest(Consumer testFunction) { + getContextRunner().run(context -> { + VectorStore vectorStore = context.getBean("vectorStore_cosine", VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test public void addAndDeleteDocumentsTest() { getContextRunner().run(context -> { @@ -148,90 +159,6 @@ class ElasticsearchVectorStoreIT { }); } - @Test - public void deleteDocumentsByFilterExpressionTest() { - getContextRunner().run(context -> { - ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine", - ElasticsearchVectorStore.class); - ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class); - - IndicesStats stats = elasticsearchClient.indices() - .stats(s -> s.index("spring-ai-document-index")) - .indices() - .get("spring-ai-document-index"); - - assertThat(stats.total().docs().count()).isEqualTo(0L); - - // Add documents with metadata - List documents = List.of( - new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")), - new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()), - new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2"))); - - vectorStore.add(documents); - elasticsearchClient.indices().refresh(); - - stats = elasticsearchClient.indices() - .stats(s -> s.index("spring-ai-document-index")) - .indices() - .get("spring-ai-document-index"); - assertThat(stats.total().docs().count()).isEqualTo(3L); - - // Delete documents with meta1 using filter expression - Expression filterExpression = new Expression(ExpressionType.EQ, new Key("meta1"), new Value("meta1")); - - vectorStore.delete(filterExpression); - elasticsearchClient.indices().refresh(); - - stats = elasticsearchClient.indices() - .stats(s -> s.index("spring-ai-document-index")) - .indices() - .get("spring-ai-document-index"); - assertThat(stats.total().docs().count()).isEqualTo(2L); - - // Clean up remaining documents - vectorStore.delete(List.of("2", "3")); - elasticsearchClient.indices().refresh(); - - stats = elasticsearchClient.indices() - .stats(s -> s.index("spring-ai-document-index")) - .indices() - .get("spring-ai-document-index"); - assertThat(stats.total().docs().count()).isEqualTo(0L); - }); - } - - @Test - public void deleteWithStringFilterExpressionTest() { - getContextRunner().run(context -> { - ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine", - ElasticsearchVectorStore.class); - ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class); - - List documents = List.of( - new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")), - new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()), - new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2"))); - - vectorStore.add(documents); - elasticsearchClient.indices().refresh(); - - // Delete documents with meta1 using string filter - vectorStore.delete("meta1 == 'meta1'"); - elasticsearchClient.indices().refresh(); - - IndicesStats stats = elasticsearchClient.indices() - .stats(s -> s.index("spring-ai-document-index")) - .indices() - .get("spring-ai-document-index"); - assertThat(stats.total().docs().count()).isEqualTo(2L); - - // Clean up - vectorStore.delete(List.of("2", "3")); - elasticsearchClient.indices().refresh(); - }); - } - @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "cosine", "l2_norm", "dot_product" }) public void addAndSearchTest(String similarityFunction) { diff --git a/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java b/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java index 315254e4b..6ad67bf69 100644 --- a/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java +++ b/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -45,6 +46,7 @@ import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -71,7 +73,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -public class MariaDBStoreIT { +public class MariaDBStoreIT extends BaseVectorStoreTests { private static String schemaName = "testdb"; @@ -141,6 +143,14 @@ public class MariaDBStoreIT { return true; } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "COSINE", "EUCLIDEAN" }) public void addAndSearch(String distanceType) { @@ -362,72 +372,6 @@ public class MariaDBStoreIT { }); } - @Test - public void deleteByFilter() { - this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - SearchRequest searchRequest = SearchRequest.builder() - .query("The World") - .topK(5) - .similarityThresholdAll() - .build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - // Verify deletion - should only have NL document remaining - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - dropTable(context); - }); - } - - @Test - public void deleteWithStringFilterExpression() { - this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - vectorStore.delete("country == 'BG'"); - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - dropTable(context); - }); - } - @Test public void deleteWithComplexFilterExpression() { this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> { diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java index 956e13819..76e81684d 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java @@ -23,6 +23,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import io.milvus.client.MilvusServiceClient; @@ -43,6 +44,7 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.embedding.TokenCountBatchingStrategy; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -64,7 +66,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -public class MilvusVectorStoreIT { +public class MilvusVectorStoreIT extends BaseVectorStoreTests { @Container private static MilvusContainer milvusContainer = new MilvusContainer(MilvusImage.DEFAULT_IMAGE); @@ -92,6 +94,15 @@ public class MilvusVectorStoreIT { ((MilvusVectorStore) vectorStore).createCollection(); } + @Override + protected void executeTest(Consumer testFunction) { + this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=" + "COSINE") + .run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "COSINE", "L2", "IP" }) public void addAndSearch(String metricType) { @@ -278,73 +289,6 @@ public class MilvusVectorStoreIT { }); } - @Test - public void deleteByFilter() { - this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - resetCollection(vectorStore); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - SearchRequest searchRequest = SearchRequest.builder() - .query("The World") - .topK(5) - .similarityThresholdAll() - .build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - // Verify deletion - should only have NL document remaining - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - - @Test - public void deleteWithStringFilterExpression() { - this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - resetCollection(vectorStore); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - // Delete using string filter expression - vectorStore.delete("country == 'BG'"); - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - @Test public void deleteWithComplexFilterExpression() { this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> { diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java index 3e98bae37..b598e4df4 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import com.mongodb.client.MongoClient; @@ -39,6 +40,7 @@ import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -64,7 +66,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -class MongoDBAtlasVectorStoreIT { +class MongoDBAtlasVectorStoreIT extends BaseVectorStoreTests { @Container private static MongoDBAtlasLocalContainer container = new MongoDBAtlasLocalContainer(MongoDbImage.DEFAULT_IMAGE); @@ -82,6 +84,14 @@ class MongoDBAtlasVectorStoreIT { }); } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test void vectorStoreTest() { this.contextRunner.run(context -> { @@ -257,71 +267,6 @@ class MongoDBAtlasVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - Thread.sleep(5000); // Wait for indexing - - SearchRequest searchRequest = SearchRequest.builder() - .query("The World") - .topK(5) - .similarityThresholdAll() - .build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - Thread.sleep(1000); // Wait for deletion to be processed - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - Thread.sleep(5000); // Wait for indexing - - var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - vectorStore.delete("country == 'BG'"); - Thread.sleep(1000); // Wait for deletion to be processed - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java index 29974c800..fbd865797 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java @@ -21,6 +21,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import org.junit.Assert; @@ -39,6 +40,7 @@ import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -60,7 +62,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -class Neo4jVectorStoreIT { +class Neo4jVectorStoreIT extends BaseVectorStoreTests { @Container static Neo4jContainer neo4jContainer = new Neo4jContainer<>(Neo4jImage.DEFAULT_IMAGE).withRandomPassword(); @@ -82,6 +84,14 @@ class Neo4jVectorStoreIT { .run(context -> context.getBean(Driver.class).executableQuery("MATCH (n) DETACH DELETE n").execute()); } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test void addAndSearchTest() { this.contextRunner.run(context -> { @@ -305,67 +315,6 @@ class Neo4jVectorStoreIT { .isTrue()); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - SearchRequest searchRequest = SearchRequest.builder() - .query("The World") - .topK(5) - .similarityThresholdAll() - .build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - vectorStore.delete("country == 'BG'"); - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java b/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java index 8eeae052b..b401983a0 100644 --- a/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java +++ b/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java @@ -416,6 +416,38 @@ class OpenSearchVectorStoreIT { }); } + @Test + void deleteById() { + getContextRunner().run(context -> { + OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class); + + var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2020, "activationDate", new Date(1000))); + var nlDocument = new Document("2", "The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "NL", "activationDate", new Date(2000))); + var bgDocument2 = new Document("3", "The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2023, "activationDate", new Date(3000))); + + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(SearchRequest.builder().query("The World").topK(5).build()), + hasSize(3)); + + vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId())); + + Awaitility.await() + .until(() -> vectorStore.similaritySearch(SearchRequest.builder().query("The World").topK(5).build()), + hasSize(1)); + + List results = vectorStore + .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); + }); + } + @Test void deleteByFilter() { getContextRunner().run(context -> { diff --git a/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java b/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java index aaa7fb9a9..854b309e7 100644 --- a/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java +++ b/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java @@ -24,6 +24,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Collectors; import javax.sql.DataSource; @@ -43,6 +44,7 @@ import org.testcontainers.utility.MountableFile; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.transformers.TransformersEmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; @@ -66,7 +68,7 @@ import static org.assertj.core.api.Assertions.assertThat; @Testcontainers @Disabled("Oracle image is 2GB") -public class OracleVectorStoreIT { +public class OracleVectorStoreIT extends BaseVectorStoreTests { @Container static OracleContainer oracle23aiContainer = new OracleContainer(OracleImage.DEFAULT_IMAGE).withCopyFileToContainer( @@ -121,6 +123,17 @@ public class OracleVectorStoreIT { return true; } + @Override + protected void executeTest(Consumer testFunction) { + this.contextRunner + .withPropertyValues("test.spring.ai.vectorstore.oracle.distanceType=COSINE", + "test.spring.ai.vectorstore.oracle.searchAccuracy=" + OracleVectorStore.DEFAULT_SEARCH_ACCURACY) + .run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "COSINE", "DOT", "EUCLIDEAN", "EUCLIDEAN_SQUARED", "MANHATTAN" }) public void addAndSearch(String distanceType) { @@ -317,71 +330,6 @@ public class OracleVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner - .withPropertyValues("test.spring.ai.vectorstore.oracle.distanceType=COSINE", - "test.spring.ai.vectorstore.oracle.searchAccuracy=" + OracleVectorStore.DEFAULT_SEARCH_ACCURACY) - .run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - List results = vectorStore.similaritySearch( - SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsKey("country") - .hasEntrySatisfying("country", - value -> assertThat(value.toString().replace("\"", "")).isEqualTo("NL")); - - dropTable(context, ((OracleVectorStore) vectorStore).getTableName()); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner - .withPropertyValues("test.spring.ai.vectorstore.oracle.distanceType=COSINE", - "test.spring.ai.vectorstore.oracle.searchAccuracy=" + OracleVectorStore.DEFAULT_SEARCH_ACCURACY) - .run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - vectorStore.delete("country == 'BG'"); - - List results = vectorStore.similaritySearch( - SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsKey("country") - .hasEntrySatisfying("country", - value -> assertThat(value.toString().replace("\"", "")).isEqualTo("NL")); - - dropTable(context, ((OracleVectorStore) vectorStore).getTableName()); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java index c80329823..112f53454 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java @@ -25,6 +25,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.function.Consumer; import java.util.stream.Stream; import javax.sql.DataSource; @@ -47,6 +48,7 @@ import org.springframework.ai.document.id.RandomIdGenerator; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.openai.OpenAiEmbeddingModel; import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -77,7 +79,7 @@ import static org.assertj.core.api.Assertions.assertThat; */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -public class PgVectorStoreIT { +public class PgVectorStoreIT extends BaseVectorStoreTests { @Container @SuppressWarnings("resource") @@ -165,6 +167,14 @@ public class PgVectorStoreIT { return true; } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "COSINE_DISTANCE", "EUCLIDEAN_DISTANCE", "NEGATIVE_INNER_PRODUCT" }) public void addAndSearch(String distanceType) { @@ -421,119 +431,6 @@ public class PgVectorStoreIT { }); } - @Test - public void deleteByIds() { - this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE") - .run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - // Create test documents - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - // Add documents to store - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - // Verify initial state - SearchRequest searchRequest = SearchRequest.builder() - .query("The World") - .topK(5) - .similarityThresholdAll() - .build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - // Delete two documents by ID - vectorStore.delete(List.of(bgDocument.getId(), nlDocument.getId())); - - // Verify deletion - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getId()).isEqualTo(bgDocument2.getId()); - - // Remove all documents from the store - dropTable(context); - }); - } - - @Test - public void deleteByFilter() { - this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE") - .run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - // Create test documents - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - // Add documents to store - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - // Verify initial state - SearchRequest searchRequest = SearchRequest.builder() - .query("The World") - .topK(5) - .similarityThresholdAll() - .build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - // Create filter to delete all documents with country=BG - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - // Delete documents using filter - vectorStore.delete(filterExpression); - - // Verify deletion - should only have NL document remaining - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - // Remove all documents from the store - dropTable(context); - }); - } - - @Test - public void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL", "year", 2021)); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build(); - - List results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(3); - - vectorStore.delete("country == 'BG'"); - - results = vectorStore.similaritySearch(searchRequest); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - vectorStore.delete(List.of(nlDocument.getId())); - }); - } - @Test void getNativeClientTest() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java b/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java index 1e183ecc4..02adc9a7c 100644 --- a/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java +++ b/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java @@ -24,6 +24,7 @@ import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; import java.util.stream.Collectors; import io.pinecone.PineconeConnection; @@ -36,6 +37,7 @@ import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.transformers.TransformersEmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; @@ -56,7 +58,7 @@ import static org.hamcrest.Matchers.hasSize; * @author Soby Chacko */ @EnabledIfEnvironmentVariable(named = "PINECONE_API_KEY", matches = ".+") -public class PineconeVectorStoreIT { +public class PineconeVectorStoreIT extends BaseVectorStoreTests { // Replace the PINECONE_ENVIRONMENT, PINECONE_PROJECT_ID, PINECONE_INDEX_NAME and // PINECONE_API_KEY with your pinecone credentials. @@ -98,6 +100,14 @@ public class PineconeVectorStoreIT { Awaitility.setDefaultTimeout(Duration.ONE_MINUTE); } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test public void addAndSearchTest() { @@ -290,59 +300,6 @@ public class PineconeVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - cleanupExistingDocuments(vectorStore, "The World"); - - var documents = createWorldDocuments(); - vectorStore.add(documents); - - awaitDocumentsCount(vectorStore, "The World", 3); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - awaitDocumentsCount(vectorStore, "The World", 1); - - List results = searchDocuments(vectorStore, "The World", 5); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - vectorStore.delete(List.of(documents.get(1).getId())); // nlDocument - awaitDocumentsCount(vectorStore, "The World", 0); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - cleanupExistingDocuments(vectorStore, "The World"); - - var documents = createWorldDocuments(); - vectorStore.add(documents); - - awaitDocumentsCount(vectorStore, "The World", 3); - - vectorStore.delete("country == 'BG'"); - - awaitDocumentsCount(vectorStore, "The World", 1); - - List results = searchDocuments(vectorStore, "The World", 5); - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - vectorStore.delete(List.of(documents.get(1).getId())); // nlDocument - awaitDocumentsCount(vectorStore, "The World", 0); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java b/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java index cf73709b3..2f45ae0cc 100644 --- a/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java +++ b/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java @@ -22,6 +22,7 @@ import java.util.Map; import java.util.Optional; import java.util.UUID; import java.util.concurrent.ExecutionException; +import java.util.function.Consumer; import java.util.stream.Collectors; import io.qdrant.client.QdrantClient; @@ -41,6 +42,7 @@ import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.mistralai.MistralAiEmbeddingModel; import org.springframework.ai.mistralai.api.MistralAiApi; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; @@ -61,7 +63,7 @@ import static org.assertj.core.api.Assertions.assertThat; @Testcontainers @EnabledIfEnvironmentVariables({ @EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+"), @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") }) -public class QdrantVectorStoreIT { +public class QdrantVectorStoreIT extends BaseVectorStoreTests { private static final String COLLECTION_NAME = "test_collection"; @@ -97,6 +99,14 @@ public class QdrantVectorStoreIT { client.close(); } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test public void addAndSearch() { this.contextRunner.run(context -> { @@ -260,57 +270,6 @@ public class QdrantVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "Bulgaria", "number", 3)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "Netherlands", "number", 90)); - - vectorStore.add(List.of(bgDocument, nlDocument)); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("Bulgaria")); - - vectorStore.delete(filterExpression); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "Netherlands"); - - vectorStore.delete(List.of(nlDocument.getId())); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "Bulgaria", "number", 3)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "Netherlands", "number", 90)); - - vectorStore.add(List.of(bgDocument, nlDocument)); - - vectorStore.delete("number > 50"); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "Bulgaria"); - - vectorStore.delete(List.of(bgDocument.getId())); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java b/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java index a2e7dd53e..a5ddf65e0 100644 --- a/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java +++ b/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java @@ -23,6 +23,8 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import java.util.stream.Collectors; import com.redis.testcontainers.RedisStackContainer; @@ -35,6 +37,7 @@ import redis.clients.jedis.JedisPooled; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.transformers.TransformersEmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; @@ -48,6 +51,7 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; import org.springframework.core.io.DefaultResourceLoader; +import org.springframework.data.redis.connection.RedisStandaloneConfiguration; import org.springframework.data.redis.connection.jedis.JedisConnectionFactory; import static org.assertj.core.api.Assertions.assertThat; @@ -59,7 +63,7 @@ import static org.assertj.core.api.Assertions.assertThat; * @author Soby Chacko */ @Testcontainers -class RedisVectorStoreIT { +class RedisVectorStoreIT extends BaseVectorStoreTests { @Container static RedisStackContainer redisContainer = new RedisStackContainer( @@ -90,6 +94,14 @@ class RedisVectorStoreIT { this.contextRunner.run(context -> context.getBean(RedisVectorStore.class).getJedis().flushAll()); } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test void ensureIndexGetsCreated() { this.contextRunner.run(context -> assertThat(context.getBean(RedisVectorStore.class).getJedis().ftList()) @@ -264,57 +276,6 @@ class RedisVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - vectorStore.delete("country == 'BG'"); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java index 0ed2fdecc..a2c0f0348 100644 --- a/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java +++ b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java @@ -25,6 +25,8 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import java.util.stream.Collectors; import org.junit.jupiter.api.Test; @@ -38,6 +40,7 @@ import org.typesense.resources.Node; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.transformers.TransformersEmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; @@ -58,7 +61,7 @@ import static org.assertj.core.api.Assertions.assertThat; * @author Thomas Vitale */ @Testcontainers -public class TypesenseVectorStoreIT { +public class TypesenseVectorStoreIT extends BaseVectorStoreTests { @Container private static TypesenseContainer typesense = new TypesenseContainer(TypesenseImage.DEFAULT_IMAGE); @@ -81,6 +84,14 @@ public class TypesenseVectorStoreIT { } } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test void documentUpdate() { this.contextRunner.run(context -> { @@ -246,61 +257,6 @@ public class TypesenseVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - ((TypesenseVectorStore) vectorStore).dropCollection(); - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - vectorStore.delete("country == 'BG'"); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - ((TypesenseVectorStore) vectorStore).dropCollection(); - }); - } - @Test void deleteWithComplexFilterExpression() { this.contextRunner.run(context -> { diff --git a/vector-stores/spring-ai-weaviate-store/pom.xml b/vector-stores/spring-ai-weaviate-store/pom.xml index 2ba8d3fa2..d603c8a0f 100644 --- a/vector-stores/spring-ai-weaviate-store/pom.xml +++ b/vector-stores/spring-ai-weaviate-store/pom.xml @@ -78,7 +78,7 @@ test - + org.springframework.boot spring-boot-starter-test test diff --git a/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java b/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java index 459d48ed8..100d42708 100644 --- a/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java +++ b/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java @@ -23,6 +23,8 @@ import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; import io.weaviate.client.Config; import io.weaviate.client.WeaviateClient; @@ -35,10 +37,10 @@ import org.testcontainers.weaviate.WeaviateContainer; import org.springframework.ai.document.Document; import org.springframework.ai.document.DocumentMetadata; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.test.vectorstore.BaseVectorStoreTests; import org.springframework.ai.transformers.TransformersEmbeddingModel; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; -import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; @@ -54,7 +56,7 @@ import static org.assertj.core.api.Assertions.assertThat; * @author Thomas Vitale */ @Testcontainers -public class WeaviateVectorStoreIT { +public class WeaviateVectorStoreIT extends BaseVectorStoreTests { @Container static WeaviateContainer weaviateContainer = new WeaviateContainer(WeaviateImage.DEFAULT_IMAGE) @@ -85,6 +87,14 @@ public class WeaviateVectorStoreIT { vectorStore.delete(this.documents.stream().map(Document::getId).toList()); } + @Override + protected void executeTest(Consumer testFunction) { + contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + testFunction.accept(vectorStore); + }); + } + @Test public void addAndSearch() { @@ -256,62 +266,6 @@ public class WeaviateVectorStoreIT { }); } - @Test - void deleteByFilter() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, - new Filter.Key("country"), new Filter.Value("BG")); - - vectorStore.delete(filterExpression); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - vectorStore.delete(List.of(nlDocument.getId())); - - }); - } - - @Test - void deleteWithStringFilterExpression() { - this.contextRunner.run(context -> { - VectorStore vectorStore = context.getBean(VectorStore.class); - - var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2020)); - var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "NL")); - var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", - Map.of("country", "BG", "year", 2023)); - - vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); - - vectorStore.delete("country == 'BG'"); - - List results = vectorStore - .similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); - - assertThat(results).hasSize(1); - assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); - - vectorStore.delete(List.of(nlDocument.getId())); - }); - } - @Test void getNativeClientTest() { this.contextRunner.run(context -> {