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 <soby.chacko@broadcom.com>
This commit is contained in:
Soby Chacko
2025-02-10 17:25:06 -05:00
committed by Ilayaperumal Gopinathan
parent a8e305d28a
commit 4d692a542b
17 changed files with 380 additions and 868 deletions

View File

@@ -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<VectorStore> testFunction);
protected Document createDocument(String country, Integer year) {
Map<String, Object> 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<Document> setupTestDocuments(VectorStore vectorStore) {
var doc1 = createDocument("BG", 2020);
var doc2 = createDocument("NL", null);
var doc3 = createDocument("BG", 2023);
List<Document> 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<Document> documents) {
await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> {
List<Document> results = vectorStore.similaritySearch(
SearchRequest.builder().query("The World").topK(documents.size()).similarityThresholdAll().build());
assertThat(results).hasSize(documents.size());
});
}
private void verifyDocumentsDeleted(VectorStore vectorStore, List<String> deletedIds) {
await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> {
List<Document> results = vectorStore
.similaritySearch(SearchRequest.builder().query("The World").topK(10).similarityThresholdAll().build());
List<String> foundIds = results.stream().map(Document::getId).collect(Collectors.toList());
assertThat(foundIds).doesNotContainAnyElementsOf(deletedIds);
});
}
@Test
protected void deleteById() {
executeTest(vectorStore -> {
List<Document> documents = setupTestDocuments(vectorStore);
verifyDocumentsExist(vectorStore, documents);
List<String> idsToDelete = List.of(documents.get(0).getId(), documents.get(1).getId());
vectorStore.delete(idsToDelete);
verifyDocumentsDeleted(vectorStore, idsToDelete);
List<Document> 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<String, Object> 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<Document> documents = setupTestDocuments(vectorStore);
verifyDocumentsExist(vectorStore, documents);
List<String> 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<Document> 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<Document> documents = setupTestDocuments(vectorStore);
verifyDocumentsExist(vectorStore, documents);
List<String> 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<Document> 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()));
});
}
}

View File

@@ -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<VectorStore> testFunction) {
contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
testFunction.accept(vectorStore);
});
}
@Override
protected Document createDocument(String country, Integer year) {
Map<String, Object> 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),

View File

@@ -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<VectorStore> 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<Document> 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() {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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) {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {

View File

@@ -78,7 +78,7 @@
<scope>test</scope>
</dependency>
<dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>

View File

@@ -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<VectorStore> 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<Document> 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<Document> 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 -> {