Skip to content

Commit d98f681

Browse files
sobychackosobychacko
committed
Add filter-based deletion to Neo4j and OpenSearch vector stores
Add string-based filter deletion alongside the Filter.Expression-based deletion for Neo4j and OpenSearch vector stores, providing consistent deletion capabilities with other vector store implementations. Key changes: - Add delete(Filter.Expression) implementation for Neo4j store using Cypher queries - Add delete(Filter.Expression) implementation for OpenSearch store using query_string - Leverage existing filter expression converters for both stores - Use Neo4j's transaction batching for efficient large-scale deletions - Use OpenSearch's delete_by_query API for metadata-based deletion - Add comprehensive integration tests for both stores covering: * Simple equality filters * String-based filter expressions * Complex filter expressions with multiple conditions This maintains consistency with other vector store implementations while utilizing store-specific features for efficient metadata-based deletion. Signed-off-by: Soby Chacko <[email protected]>
1 parent 8403049 commit d98f681

File tree

4 files changed

+298
-6
lines changed

4 files changed

+298
-6
lines changed

‎vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,18 @@
2525
import org.neo4j.driver.Driver;
2626
import org.neo4j.driver.SessionConfig;
2727
import org.neo4j.driver.Values;
28+
import org.slf4j.Logger;
29+
import org.slf4j.LoggerFactory;
2830

2931
import org.springframework.ai.document.Document;
3032
import org.springframework.ai.document.DocumentMetadata;
31-
import org.springframework.ai.embedding.BatchingStrategy;
3233
import org.springframework.ai.embedding.EmbeddingModel;
3334
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
34-
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
3535
import org.springframework.ai.observation.conventions.VectorStoreProvider;
3636
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
3737
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
3838
import org.springframework.ai.vectorstore.SearchRequest;
39+
import org.springframework.ai.vectorstore.filter.Filter;
3940
import org.springframework.ai.vectorstore.neo4j.filter.Neo4jVectorFilterExpressionConverter;
4041
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
4142
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
@@ -133,6 +134,8 @@
133134
*/
134135
public class Neo4jVectorStore extends AbstractObservationVectorStore implements InitializingBean {
135136

137+
private static final Logger logger = LoggerFactory.getLogger(Neo4jVectorStore.class);
138+
136139
public static final int DEFAULT_EMBEDDING_DIMENSION = 1536;
137140

138141
public static final int DEFAULT_TRANSACTION_SIZE = 10_000;
@@ -235,6 +238,31 @@ public Optional<Boolean> doDelete(List<String> idList) {
235238
}
236239
}
237240

241+
@Override
242+
protected void doDelete(Filter.Expression filterExpression) {
243+
Assert.notNull(filterExpression, "Filter expression must not be null");
244+
245+
try (var session = this.driver.session(this.sessionConfig)) {
246+
String whereClause = this.filterExpressionConverter.convertExpression(filterExpression);
247+
248+
// Create Cypher query with transaction batching
249+
String cypher = """
250+
MATCH (node:%s) WHERE %s
251+
CALL { WITH node DETACH DELETE node } IN TRANSACTIONS OF $transactionSize ROWS
252+
""".formatted(this.label, whereClause);
253+
254+
var summary = session.run(cypher,
255+
Map.of("transactionSize", DEFAULT_TRANSACTION_SIZE))
256+
.consume();
257+
258+
logger.debug("Deleted {} nodes matching filter expression", summary.counters().nodesDeleted());
259+
}
260+
catch (Exception e) {
261+
logger.error("Failed to delete nodes by filter: {}", e.getMessage(), e);
262+
throw new IllegalStateException("Failed to delete nodes by filter", e);
263+
}
264+
}
265+
238266
@Override
239267
public List<Document> doSimilaritySearch(SearchRequest request) {
240268
Assert.isTrue(request.getTopK() > 0, "The number of documents to returned must be greater than zero");

‎vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java‎

Lines changed: 108 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -20,6 +20,7 @@
2020
import java.util.List;
2121
import java.util.Map;
2222
import java.util.UUID;
23+
import java.util.stream.Collectors;
2324

2425
import org.junit.Assert;
2526
import org.junit.jupiter.api.BeforeEach;
@@ -39,6 +40,7 @@
3940
import org.springframework.ai.openai.api.OpenAiApi;
4041
import org.springframework.ai.vectorstore.SearchRequest;
4142
import org.springframework.ai.vectorstore.VectorStore;
43+
import org.springframework.ai.vectorstore.filter.Filter;
4244
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
4345
import org.springframework.boot.SpringBootConfiguration;
4446
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -53,6 +55,7 @@
5355
* @author Michael Simons
5456
* @author Christian Tzolov
5557
* @author Thomas Vitale
58+
* @author Soby Chacko
5659
*/
5760
@Testcontainers
5861
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -301,6 +304,110 @@ void ensureIdIndexGetsCreated() {
301304
.isTrue());
302305
}
303306

307+
@Test
308+
void deleteByFilter() {
309+
this.contextRunner.run(context -> {
310+
VectorStore vectorStore = context.getBean(VectorStore.class);
311+
312+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
313+
Map.of("country", "BG", "year", 2020));
314+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
315+
Map.of("country", "NL", "year", 2021));
316+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
317+
Map.of("country", "BG", "year", 2023));
318+
319+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
320+
321+
SearchRequest searchRequest = SearchRequest.builder()
322+
.query("The World")
323+
.topK(5)
324+
.similarityThresholdAll()
325+
.build();
326+
327+
List<Document> results = vectorStore.similaritySearch(searchRequest);
328+
assertThat(results).hasSize(3);
329+
330+
Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
331+
new Filter.Key("country"), new Filter.Value("BG"));
332+
333+
vectorStore.delete(filterExpression);
334+
335+
results = vectorStore.similaritySearch(searchRequest);
336+
assertThat(results).hasSize(1);
337+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
338+
});
339+
}
340+
341+
@Test
342+
void deleteWithStringFilterExpression() {
343+
this.contextRunner.run(context -> {
344+
VectorStore vectorStore = context.getBean(VectorStore.class);
345+
346+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
347+
Map.of("country", "BG", "year", 2020));
348+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
349+
Map.of("country", "NL", "year", 2021));
350+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
351+
Map.of("country", "BG", "year", 2023));
352+
353+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
354+
355+
var searchRequest = SearchRequest.builder()
356+
.query("The World")
357+
.topK(5)
358+
.similarityThresholdAll()
359+
.build();
360+
361+
List<Document> results = vectorStore.similaritySearch(searchRequest);
362+
assertThat(results).hasSize(3);
363+
364+
vectorStore.delete("country == 'BG'");
365+
366+
results = vectorStore.similaritySearch(searchRequest);
367+
assertThat(results).hasSize(1);
368+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
369+
});
370+
}
371+
372+
@Test
373+
void deleteWithComplexFilterExpression() {
374+
this.contextRunner.run(context -> {
375+
VectorStore vectorStore = context.getBean(VectorStore.class);
376+
377+
var doc1 = new Document("Content 1", Map.of("type", "A", "priority", 1L));
378+
var doc2 = new Document("Content 2", Map.of("type", "A", "priority", 2L));
379+
var doc3 = new Document("Content 3", Map.of("type", "B", "priority", 1L));
380+
381+
vectorStore.add(List.of(doc1, doc2, doc3));
382+
383+
// Complex filter expression: (type == 'A' AND priority > 1)
384+
Filter.Expression priorityFilter = new Filter.Expression(Filter.ExpressionType.GT,
385+
new Filter.Key("priority"), new Filter.Value(1));
386+
Filter.Expression typeFilter = new Filter.Expression(Filter.ExpressionType.EQ,
387+
new Filter.Key("type"), new Filter.Value("A"));
388+
Filter.Expression complexFilter = new Filter.Expression(Filter.ExpressionType.AND,
389+
typeFilter, priorityFilter);
390+
391+
vectorStore.delete(complexFilter);
392+
393+
var results = vectorStore.similaritySearch(SearchRequest.builder()
394+
.query("Content")
395+
.topK(5)
396+
.similarityThresholdAll()
397+
.build());
398+
399+
assertThat(results).hasSize(2);
400+
assertThat(results.stream()
401+
.map(doc -> doc.getMetadata().get("type"))
402+
.collect(Collectors.toList()))
403+
.containsExactlyInAnyOrder("A", "B");
404+
assertThat(results.stream()
405+
.map(doc -> doc.getMetadata().get("priority"))
406+
.collect(Collectors.toList()))
407+
.containsExactlyInAnyOrder(1L, 1L);
408+
});
409+
}
410+
304411
@SpringBootConfiguration
305412
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
306413
public static class TestApplication {

‎vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,17 +32,19 @@
3232
import org.opensearch.client.opensearch._types.query_dsl.Query;
3333
import org.opensearch.client.opensearch.core.BulkRequest;
3434
import org.opensearch.client.opensearch.core.BulkResponse;
35+
import org.opensearch.client.opensearch.core.DeleteByQueryRequest;
36+
import org.opensearch.client.opensearch.core.DeleteByQueryResponse;
3537
import org.opensearch.client.opensearch.core.search.Hit;
3638
import org.opensearch.client.opensearch.indices.CreateIndexRequest;
3739
import org.opensearch.client.opensearch.indices.CreateIndexResponse;
3840
import org.opensearch.client.transport.endpoints.BooleanResponse;
41+
import org.slf4j.Logger;
42+
import org.slf4j.LoggerFactory;
3943

4044
import org.springframework.ai.document.Document;
4145
import org.springframework.ai.document.DocumentMetadata;
42-
import org.springframework.ai.embedding.BatchingStrategy;
4346
import org.springframework.ai.embedding.EmbeddingModel;
4447
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
45-
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
4648
import org.springframework.ai.observation.conventions.VectorStoreProvider;
4749
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
4850
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
@@ -139,6 +141,8 @@
139141
*/
140142
public class OpenSearchVectorStore extends AbstractObservationVectorStore implements InitializingBean {
141143

144+
private static final Logger logger = LoggerFactory.getLogger(OpenSearchVectorStore.class);
145+
142146
public static final String COSINE_SIMILARITY_FUNCTION = "cosinesimil";
143147

144148
public static final String DEFAULT_INDEX_NAME = "spring-ai-document-index";
@@ -230,6 +234,35 @@ private BulkResponse bulkRequest(BulkRequest bulkRequest) {
230234
}
231235
}
232236

237+
@Override
238+
protected void doDelete(Filter.Expression filterExpression) {
239+
Assert.notNull(filterExpression, "Filter expression must not be null");
240+
241+
try {
242+
String filterStr = this.filterExpressionConverter.convertExpression(filterExpression);
243+
244+
// Create delete by query request
245+
DeleteByQueryRequest request = new DeleteByQueryRequest.Builder()
246+
.index(this.index)
247+
.query(q -> q
248+
.queryString(qs -> qs
249+
.query(filterStr)
250+
))
251+
.build();
252+
253+
DeleteByQueryResponse response = this.openSearchClient.deleteByQuery(request);
254+
logger.debug("Deleted {} documents matching filter expression", response.deleted());
255+
256+
if (!response.failures().isEmpty()) {
257+
throw new IllegalStateException("Failed to delete some documents: " + response.failures());
258+
}
259+
}
260+
catch (Exception e) {
261+
logger.error("Failed to delete documents by filter: {}", e.getMessage(), e);
262+
throw new IllegalStateException("Failed to delete documents by filter", e);
263+
}
264+
}
265+
233266
@Override
234267
public List<Document> doSimilaritySearch(SearchRequest searchRequest) {
235268
Assert.notNull(searchRequest, "The search request must not be null.");

0 commit comments

Comments
 (0)