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:
committed by
Ilayaperumal Gopinathan
parent
a8e305d28a
commit
4d692a542b
@@ -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()));
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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() {
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
@@ -78,7 +78,7 @@
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-test</artifactId>
|
||||
<scope>test</scope>
|
||||
|
||||
@@ -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 -> {
|
||||
|
||||
Reference in New Issue
Block a user