Add filter-based deletion to MongoDB Atlas vector store
Add string-based filter deletion alongside the Filter.Expression-based deletion for MongoDB Atlas vector store, providing consistent deletion capabilities with other vector store implementations. Key changes: - Add delete(Filter.Expression) implementation for MongoDB Atlas store - Leverage existing MongoDBAtlasFilterExpressionConverter for filter translation - Use MongoTemplate's native query capabilities for deletion - Add comprehensive integration tests for filter deletion - Support both simple and complex filter expressions This maintains consistency with other vector store implementations while utilizing MongoDB-specific query capabilities for efficient metadata-based deletion.
This commit is contained in:
committed by
Mark Pollack
parent
c5dd7683b6
commit
ef0bade409
@@ -23,23 +23,25 @@ import java.util.Map;
|
||||
import java.util.Optional;
|
||||
|
||||
import com.mongodb.MongoCommandException;
|
||||
import com.mongodb.client.result.DeleteResult;
|
||||
import org.springframework.core.log.LogAccessor;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentMetadata;
|
||||
import org.springframework.ai.embedding.BatchingStrategy;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
|
||||
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
|
||||
import org.springframework.ai.model.EmbeddingUtils;
|
||||
import org.springframework.ai.observation.conventions.VectorStoreProvider;
|
||||
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
|
||||
import org.springframework.beans.factory.InitializingBean;
|
||||
import org.springframework.data.mongodb.UncategorizedMongoDbException;
|
||||
import org.springframework.data.mongodb.core.MongoTemplate;
|
||||
import org.springframework.data.mongodb.core.aggregation.Aggregation;
|
||||
import org.springframework.data.mongodb.core.query.BasicQuery;
|
||||
import org.springframework.data.mongodb.core.query.Criteria;
|
||||
import org.springframework.data.mongodb.core.query.Query;
|
||||
import org.springframework.util.Assert;
|
||||
@@ -125,6 +127,8 @@ import org.springframework.util.Assert;
|
||||
*/
|
||||
public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore implements InitializingBean {
|
||||
|
||||
private static final LogAccessor logger = new LogAccessor(MongoDBAtlasVectorStore.class);
|
||||
|
||||
public static final String ID_FIELD_NAME = "_id";
|
||||
|
||||
public static final String METADATA_FIELD_NAME = "metadata";
|
||||
@@ -272,6 +276,22 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
|
||||
return Optional.of(deleteCount == idList.size());
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void doDelete(Filter.Expression filterExpression) {
|
||||
Assert.notNull(filterExpression, "Filter expression must not be null");
|
||||
|
||||
try {
|
||||
String nativeFilterExpression = this.filterExpressionConverter.convertExpression(filterExpression);
|
||||
BasicQuery query = new BasicQuery(nativeFilterExpression);
|
||||
DeleteResult deleteResult = this.mongoTemplate.remove(query, this.collectionName);
|
||||
|
||||
logger.debug("Deleted " + deleteResult.getDeletedCount() + " documents matching filter expression");
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new IllegalStateException("Failed to delete documents by filter", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> similaritySearch(String query) {
|
||||
return similaritySearch(SearchRequest.builder().query(query).build());
|
||||
|
||||
@@ -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.
|
||||
@@ -40,6 +40,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.autoconfigure.EnableAutoConfiguration;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
@@ -255,6 +256,105 @@ 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 -> {
|
||||
VectorStore vectorStore = context.getBean(VectorStore.class);
|
||||
|
||||
var doc1 = new Document("Content 1", Map.of("type", "A", "priority", 1));
|
||||
var doc2 = new Document("Content 2", Map.of("type", "A", "priority", 2));
|
||||
var doc3 = new Document("Content 3", Map.of("type", "B", "priority", 1));
|
||||
|
||||
vectorStore.add(List.of(doc1, doc2, doc3));
|
||||
Thread.sleep(5000); // Wait for indexing
|
||||
|
||||
// Complex filter expression: (type == 'A' AND priority > 1)
|
||||
Filter.Expression priorityFilter = new Filter.Expression(Filter.ExpressionType.GT,
|
||||
new Filter.Key("priority"), new Filter.Value(1));
|
||||
Filter.Expression typeFilter = new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("type"),
|
||||
new Filter.Value("A"));
|
||||
Filter.Expression complexFilter = new Filter.Expression(Filter.ExpressionType.AND, typeFilter,
|
||||
priorityFilter);
|
||||
|
||||
vectorStore.delete(complexFilter);
|
||||
Thread.sleep(1000); // Wait for deletion to be processed
|
||||
|
||||
var results = vectorStore
|
||||
.similaritySearch(SearchRequest.builder().query("Content").topK(5).similarityThresholdAll().build());
|
||||
|
||||
assertThat(results).hasSize(2);
|
||||
assertThat(results.stream().map(doc -> doc.getMetadata().get("type")).collect(Collectors.toList()))
|
||||
.containsExactlyInAnyOrder("A", "B");
|
||||
assertThat(results.stream().map(doc -> doc.getMetadata().get("priority")).collect(Collectors.toList()))
|
||||
.containsExactlyInAnyOrder(1, 1);
|
||||
});
|
||||
}
|
||||
|
||||
public static String getText(String uri) {
|
||||
var resource = new DefaultResourceLoader().getResource(uri);
|
||||
try {
|
||||
|
||||
Reference in New Issue
Block a user