Add filter-based deletion and string filter API to VectorStore
Add string-based filter deletion alongside the Filter.Expression-based deletion for vector stores. This provides a more convenient API for simple filter cases while maintaining the full flexibility of expression-based filtering. Key changes: - Add delete(Filter.Expression) and delete(String) methods to VectorStore - Add default string filter implementation in VectorStore interface - Implement filter deletion in Chroma/Elasticsearch/PgVector stores - Add integration tests for both filter APIs across implementations This extends vector store deletion capabilities while maintaining consistent behavior across implementations.
This commit is contained in:
committed by
Mark Pollack
parent
7521cab9db
commit
f588716b41
@@ -1,5 +1,5 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
* 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.
|
||||
@@ -24,9 +24,11 @@ import io.micrometer.observation.ObservationRegistry;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentWriter;
|
||||
import org.springframework.ai.embedding.BatchingStrategy;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
|
||||
import org.springframework.lang.Nullable;
|
||||
import org.springframework.util.Assert;
|
||||
|
||||
/**
|
||||
* The {@code VectorStore} interface defines the operations for managing and querying
|
||||
@@ -62,6 +64,28 @@ public interface VectorStore extends DocumentWriter {
|
||||
@Nullable
|
||||
Optional<Boolean> delete(List<String> idList);
|
||||
|
||||
/**
|
||||
* Deletes documents from the vector store based on filter criteria.
|
||||
* @param filterExpression Filter expression to identify documents to delete
|
||||
* @throws IllegalStateException if the underlying delete causes an exception
|
||||
*/
|
||||
void delete(Filter.Expression filterExpression);
|
||||
|
||||
/**
|
||||
* Deletes documents from the vector store using a string filter expression. Converts
|
||||
* the string filter to an Expression object and delegates to
|
||||
* {@link #delete(Filter.Expression)}.
|
||||
* @param filterExpression String representation of the filter criteria
|
||||
* @throws IllegalArgumentException if the filter expression is null
|
||||
* @throws IllegalStateException if the underlying delete causes an exception
|
||||
*/
|
||||
default void delete(String filterExpression) {
|
||||
SearchRequest searchRequest = SearchRequest.builder().filterExpression(filterExpression).build();
|
||||
Filter.Expression textExpression = searchRequest.getFilterExpression();
|
||||
Assert.notNull(textExpression, "Filter expression must not be null");
|
||||
this.delete(textExpression);
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves documents by query embedding similarity and metadata filters to retrieve
|
||||
* exactly the number of nearest-neighbor results that match the request criteria.
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
* 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.
|
||||
@@ -18,6 +18,7 @@ package org.springframework.ai.vectorstore.observation;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import io.micrometer.observation.ObservationRegistry;
|
||||
|
||||
@@ -27,6 +28,7 @@ import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.lang.Nullable;
|
||||
|
||||
/**
|
||||
@@ -99,6 +101,18 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
|
||||
.observe(() -> this.doDelete(deleteDocIds));
|
||||
}
|
||||
|
||||
@Override
|
||||
public void delete(Filter.Expression filterExpression) {
|
||||
VectorStoreObservationContext observationContext = this
|
||||
.createObservationContextBuilder(VectorStoreObservationContext.Operation.DELETE.value())
|
||||
.build();
|
||||
|
||||
VectorStoreObservationDocumentation.AI_VECTOR_STORE
|
||||
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
|
||||
this.observationRegistry)
|
||||
.observe(() -> this.doDelete(filterExpression));
|
||||
}
|
||||
|
||||
@Override
|
||||
@Nullable
|
||||
public List<Document> similaritySearch(SearchRequest request) {
|
||||
@@ -131,6 +145,19 @@ public abstract class AbstractObservationVectorStore implements VectorStore {
|
||||
*/
|
||||
public abstract Optional<Boolean> doDelete(List<String> idList);
|
||||
|
||||
/**
|
||||
* Template method for concrete implementations to provide filter-based deletion
|
||||
* logic.
|
||||
* @param filterExpression Filter expression to identify documents to delete
|
||||
*/
|
||||
@Nullable
|
||||
protected void doDelete(Filter.Expression filterExpression) {
|
||||
// this is temporary until we implement this method in all concrete vector stores,
|
||||
// at which point
|
||||
// this method will become an abstract method.
|
||||
throw new UnsupportedOperationException();
|
||||
}
|
||||
|
||||
/**
|
||||
* Perform the actual similarity search operation.
|
||||
* @param request the search request
|
||||
|
||||
@@ -27,12 +27,20 @@ Spring AI offers an abstracted API for interacting with vector databases through
|
||||
Here is the `VectorStore` interface definition:
|
||||
|
||||
```java
|
||||
public interface VectorStore {
|
||||
public interface VectorStore extends DocumentWriter {
|
||||
|
||||
default String getName() {
|
||||
return this.getClass().getSimpleName();
|
||||
}
|
||||
|
||||
void add(List<Document> documents);
|
||||
|
||||
Optional<Boolean> delete(List<String> idList);
|
||||
|
||||
void delete(Filter.Expression filterExpression);
|
||||
|
||||
default void delete(String filterExpression) { ... };
|
||||
|
||||
List<Document> similaritySearch(String query);
|
||||
|
||||
List<Document> similaritySearch(SearchRequest request);
|
||||
|
||||
@@ -338,7 +338,7 @@ public class ChromaApi {
|
||||
*/
|
||||
@JsonInclude(JsonInclude.Include.NON_NULL)
|
||||
public record DeleteEmbeddingsRequest(// @formatter:off
|
||||
@JsonProperty("ids") List<String> ids,
|
||||
@Nullable @JsonProperty("ids") List<String> ids,
|
||||
@Nullable @JsonProperty("where") Map<String, Object> where) { // @formatter:on
|
||||
|
||||
public DeleteEmbeddingsRequest(List<String> ids) {
|
||||
|
||||
@@ -25,6 +25,8 @@ import java.util.Optional;
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.fasterxml.jackson.databind.json.JsonMapper;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import org.springframework.ai.chroma.vectorstore.ChromaApi.AddEmbeddingsRequest;
|
||||
import org.springframework.ai.chroma.vectorstore.ChromaApi.DeleteEmbeddingsRequest;
|
||||
@@ -40,6 +42,7 @@ import org.springframework.ai.util.JacksonUtils;
|
||||
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
|
||||
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
|
||||
@@ -81,6 +84,8 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
|
||||
|
||||
private boolean initialized = false;
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(ChromaVectorStore.class);
|
||||
|
||||
/**
|
||||
* @param builder {@link VectorStore.Builder} for chroma vector store
|
||||
*/
|
||||
@@ -162,7 +167,29 @@ public class ChromaVectorStore extends AbstractObservationVectorStore implements
|
||||
}
|
||||
|
||||
@Override
|
||||
public @NonNull List<Document> doSimilaritySearch(@NonNull SearchRequest request) {
|
||||
protected void doDelete(Filter.Expression expression) {
|
||||
Assert.notNull(expression, "Filter expression must not be null");
|
||||
|
||||
try {
|
||||
ChromaFilterExpressionConverter converter = new ChromaFilterExpressionConverter();
|
||||
String whereClauseStr = converter.convertExpression(expression);
|
||||
|
||||
Map<String, Object> whereClause = this.chromaApi.where(whereClauseStr);
|
||||
|
||||
logger.debug("Deleting with where clause: {}", whereClause);
|
||||
|
||||
DeleteEmbeddingsRequest deleteRequest = new DeleteEmbeddingsRequest(null, whereClause);
|
||||
this.chromaApi.deleteEmbeddings(this.collectionId, deleteRequest);
|
||||
}
|
||||
catch (Exception e) {
|
||||
logger.error("Failed to delete documents by filter: {}", e.getMessage(), e);
|
||||
throw new IllegalStateException("Failed to delete documents by filter", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
@NonNull
|
||||
public List<Document> doSimilaritySearch(@NonNull SearchRequest request) {
|
||||
|
||||
String query = request.getQuery();
|
||||
Assert.notNull(query, "Query string must not be null");
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
* 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.
|
||||
@@ -36,6 +36,7 @@ import org.springframework.ai.openai.OpenAiEmbeddingModel;
|
||||
import org.springframework.ai.openai.api.OpenAiApi;
|
||||
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.test.context.runner.ApplicationContextRunner;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
@@ -169,6 +170,69 @@ 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() {
|
||||
|
||||
|
||||
@@ -26,6 +26,7 @@ import java.util.stream.Collectors;
|
||||
import co.elastic.clients.elasticsearch.ElasticsearchClient;
|
||||
import co.elastic.clients.elasticsearch.core.BulkRequest;
|
||||
import co.elastic.clients.elasticsearch.core.BulkResponse;
|
||||
import co.elastic.clients.elasticsearch.core.DeleteByQueryResponse;
|
||||
import co.elastic.clients.elasticsearch.core.SearchResponse;
|
||||
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
|
||||
import co.elastic.clients.elasticsearch.core.search.Hit;
|
||||
@@ -221,6 +222,23 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
|
||||
return Optional.of(bulkRequest(bulkRequestBuilder.build()).errors());
|
||||
}
|
||||
|
||||
@Override
|
||||
public void doDelete(Filter.Expression filterExpression) {
|
||||
// For the index to be present, either it must be pre-created or set the
|
||||
// initializeSchema to true.
|
||||
if (!indexExists()) {
|
||||
throw new IllegalArgumentException("Index not found");
|
||||
}
|
||||
|
||||
try {
|
||||
this.elasticsearchClient.deleteByQuery(d -> d.index(this.options.getIndexName())
|
||||
.query(q -> q.queryString(qs -> qs.query(getElasticsearchQueryString(filterExpression)))));
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new IllegalStateException("Failed to delete documents by filter", e);
|
||||
}
|
||||
}
|
||||
|
||||
private BulkResponse bulkRequest(BulkRequest bulkRequest) {
|
||||
try {
|
||||
return this.elasticsearchClient.bulk(bulkRequest);
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
* 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.
|
||||
@@ -53,6 +53,11 @@ import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.openai.OpenAiEmbeddingModel;
|
||||
import org.springframework.ai.openai.api.OpenAiApi;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.Filter.Expression;
|
||||
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
|
||||
import org.springframework.ai.vectorstore.filter.Filter.Key;
|
||||
import org.springframework.ai.vectorstore.filter.Filter.Value;
|
||||
import org.springframework.boot.SpringBootConfiguration;
|
||||
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
|
||||
@@ -143,6 +148,90 @@ 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) {
|
||||
|
||||
@@ -45,6 +45,7 @@ import org.springframework.ai.util.JacksonUtils;
|
||||
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
|
||||
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
|
||||
@@ -315,6 +316,22 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
|
||||
return Optional.of(updateCount == idList.size());
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doDelete(Filter.Expression filterExpression) {
|
||||
String nativeFilterExpression = this.filterExpressionConverter.convertExpression(filterExpression);
|
||||
|
||||
String sql = "DELETE FROM " + getFullyQualifiedTableName() + " WHERE metadata::jsonb @@ '"
|
||||
+ nativeFilterExpression + "'::jsonpath";
|
||||
|
||||
// Execute the delete
|
||||
try {
|
||||
this.jdbcTemplate.update(sql);
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new IllegalStateException("Failed to delete documents by filter", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> doSimilaritySearch(SearchRequest request) {
|
||||
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
/*
|
||||
* Copyright 2023-2024 the original author or authors.
|
||||
* 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.
|
||||
@@ -22,6 +22,7 @@ import java.util.Collections;
|
||||
import java.util.Iterator;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.UUID;
|
||||
import java.util.stream.Stream;
|
||||
|
||||
@@ -29,6 +30,7 @@ import javax.sql.DataSource;
|
||||
|
||||
import com.zaxxer.hikari.HikariDataSource;
|
||||
import org.junit.Assert;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.Arguments;
|
||||
@@ -45,6 +47,7 @@ import org.springframework.ai.openai.OpenAiEmbeddingModel;
|
||||
import org.springframework.ai.openai.api.OpenAiApi;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
|
||||
import org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
@@ -364,6 +367,121 @@ 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
|
||||
Optional<Boolean> deleteResult = vectorStore.delete(List.of(bgDocument.getId(), nlDocument.getId()));
|
||||
|
||||
assertThat(deleteResult).isPresent().contains(true);
|
||||
|
||||
// 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()));
|
||||
});
|
||||
}
|
||||
|
||||
@SpringBootConfiguration
|
||||
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
|
||||
public static class TestApplication {
|
||||
|
||||
Reference in New Issue
Block a user