|
1 | 1 | /*
|
2 |
| - * Copyright 2023-2024 the original author or authors. |
| 2 | + * Copyright 2023-2025 the original author or authors. |
3 | 3 | *
|
4 | 4 | * Licensed under the Apache License, Version 2.0 (the "License");
|
5 | 5 | * you may not use this file except in compliance with the License.
|
|
23 | 23 | import java.util.List;
|
24 | 24 | import java.util.Map;
|
25 | 25 | import java.util.UUID;
|
| 26 | +import java.util.stream.Collectors; |
26 | 27 | import java.util.stream.Stream;
|
27 | 28 |
|
28 | 29 | import javax.sql.DataSource;
|
29 | 30 |
|
30 | 31 | import com.zaxxer.hikari.HikariDataSource;
|
31 | 32 | import org.junit.Assert;
|
| 33 | +import org.junit.jupiter.api.Test; |
32 | 34 | import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
33 | 35 | import org.junit.jupiter.params.ParameterizedTest;
|
34 | 36 | import org.junit.jupiter.params.provider.Arguments;
|
|
44 | 46 | import org.springframework.ai.openai.api.OpenAiApi;
|
45 | 47 | import org.springframework.ai.vectorstore.SearchRequest;
|
46 | 48 | import org.springframework.ai.vectorstore.VectorStore;
|
| 49 | +import org.springframework.ai.vectorstore.filter.Filter; |
47 | 50 | import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
|
48 | 51 | import org.springframework.beans.factory.annotation.Value;
|
49 | 52 | import org.springframework.boot.SpringBootConfiguration;
|
|
63 | 66 |
|
64 | 67 | /**
|
65 | 68 | * @author Diego Dupin
|
| 69 | + * @author Soby Chacko |
66 | 70 | */
|
67 | 71 | @Testcontainers
|
68 | 72 | @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
|
@@ -357,6 +361,106 @@ public void searchWithThreshold(String distanceType) {
|
357 | 361 | });
|
358 | 362 | }
|
359 | 363 |
|
| 364 | + @Test |
| 365 | + public void deleteByFilter() { |
| 366 | + this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> { |
| 367 | + VectorStore vectorStore = context.getBean(VectorStore.class); |
| 368 | + |
| 369 | + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 370 | + Map.of("country", "BG", "year", 2020)); |
| 371 | + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 372 | + Map.of("country", "NL", "year", 2021)); |
| 373 | + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 374 | + Map.of("country", "BG", "year", 2023)); |
| 375 | + |
| 376 | + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); |
| 377 | + |
| 378 | + SearchRequest searchRequest = SearchRequest.builder() |
| 379 | + .query("The World") |
| 380 | + .topK(5) |
| 381 | + .similarityThresholdAll() |
| 382 | + .build(); |
| 383 | + |
| 384 | + List<Document> results = vectorStore.similaritySearch(searchRequest); |
| 385 | + assertThat(results).hasSize(3); |
| 386 | + |
| 387 | + Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, |
| 388 | + new Filter.Key("country"), new Filter.Value("BG")); |
| 389 | + |
| 390 | + vectorStore.delete(filterExpression); |
| 391 | + |
| 392 | + // Verify deletion - should only have NL document remaining |
| 393 | + results = vectorStore.similaritySearch(searchRequest); |
| 394 | + assertThat(results).hasSize(1); |
| 395 | + assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); |
| 396 | + |
| 397 | + dropTable(context); |
| 398 | + }); |
| 399 | + } |
| 400 | + |
| 401 | + @Test |
| 402 | + public void deleteWithStringFilterExpression() { |
| 403 | + this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> { |
| 404 | + VectorStore vectorStore = context.getBean(VectorStore.class); |
| 405 | + |
| 406 | + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 407 | + Map.of("country", "BG", "year", 2020)); |
| 408 | + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 409 | + Map.of("country", "NL", "year", 2021)); |
| 410 | + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 411 | + Map.of("country", "BG", "year", 2023)); |
| 412 | + |
| 413 | + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); |
| 414 | + |
| 415 | + var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build(); |
| 416 | + |
| 417 | + List<Document> results = vectorStore.similaritySearch(searchRequest); |
| 418 | + assertThat(results).hasSize(3); |
| 419 | + |
| 420 | + vectorStore.delete("country == 'BG'"); |
| 421 | + |
| 422 | + results = vectorStore.similaritySearch(searchRequest); |
| 423 | + assertThat(results).hasSize(1); |
| 424 | + assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); |
| 425 | + |
| 426 | + dropTable(context); |
| 427 | + }); |
| 428 | + } |
| 429 | + |
| 430 | + @Test |
| 431 | + public void deleteWithComplexFilterExpression() { |
| 432 | + this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> { |
| 433 | + VectorStore vectorStore = context.getBean(VectorStore.class); |
| 434 | + |
| 435 | + var doc1 = new Document("Content 1", Map.of("type", "A", "priority", 1)); |
| 436 | + var doc2 = new Document("Content 2", Map.of("type", "A", "priority", 2)); |
| 437 | + var doc3 = new Document("Content 3", Map.of("type", "B", "priority", 1)); |
| 438 | + |
| 439 | + vectorStore.add(List.of(doc1, doc2, doc3)); |
| 440 | + |
| 441 | + // Complex filter expression: (type == 'A' AND priority > 1) |
| 442 | + Filter.Expression priorityFilter = new Filter.Expression(Filter.ExpressionType.GT, |
| 443 | + new Filter.Key("priority"), new Filter.Value(1)); |
| 444 | + Filter.Expression typeFilter = new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("type"), |
| 445 | + new Filter.Value("A")); |
| 446 | + Filter.Expression complexFilter = new Filter.Expression(Filter.ExpressionType.AND, typeFilter, |
| 447 | + priorityFilter); |
| 448 | + |
| 449 | + vectorStore.delete(complexFilter); |
| 450 | + |
| 451 | + var results = vectorStore |
| 452 | + .similaritySearch(SearchRequest.builder().query("Content").topK(5).similarityThresholdAll().build()); |
| 453 | + |
| 454 | + assertThat(results).hasSize(2); |
| 455 | + assertThat(results.stream().map(doc -> doc.getMetadata().get("type")).collect(Collectors.toList())) |
| 456 | + .containsExactlyInAnyOrder("A", "B"); |
| 457 | + assertThat(results.stream().map(doc -> doc.getMetadata().get("priority")).collect(Collectors.toList())) |
| 458 | + .containsExactlyInAnyOrder(1, 1); |
| 459 | + |
| 460 | + dropTable(context); |
| 461 | + }); |
| 462 | + } |
| 463 | + |
360 | 464 | @SpringBootConfiguration
|
361 | 465 | @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
|
362 | 466 | public static class TestApplication {
|
|
0 commit comments