|
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.
|
|
22 | 22 | import java.util.List;
|
23 | 23 | import java.util.Map;
|
24 | 24 | import java.util.UUID;
|
| 25 | +import java.util.stream.Collectors; |
25 | 26 |
|
26 | 27 | import com.datastax.oss.driver.api.core.CqlSession;
|
27 | 28 | import com.datastax.oss.driver.api.core.CqlSessionBuilder;
|
|
42 | 43 | import org.springframework.ai.vectorstore.SearchRequest;
|
43 | 44 | import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn;
|
44 | 45 | import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumnTags;
|
| 46 | +import org.springframework.ai.vectorstore.filter.Filter; |
45 | 47 | import org.springframework.boot.SpringBootConfiguration;
|
46 | 48 | import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
|
47 | 49 | import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
|
|
57 | 59 | *
|
58 | 60 | * @author Mick Semb Wever
|
59 | 61 | * @author Thomas Vitale
|
| 62 | + * @author Soby Chacko |
60 | 63 | * @since 1.0.0
|
61 | 64 | */
|
62 | 65 | @Testcontainers
|
@@ -417,6 +420,110 @@ void searchWithThreshold() {
|
417 | 420 | });
|
418 | 421 | }
|
419 | 422 |
|
| 423 | + @Test |
| 424 | + void deleteByFilter() { |
| 425 | + this.contextRunner.run(context -> { |
| 426 | + try (CassandraVectorStore store = createTestStore(context, |
| 427 | + new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED), |
| 428 | + new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) { |
| 429 | + |
| 430 | + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 431 | + Map.of("country", "BG", "year", (short) 2020)); |
| 432 | + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 433 | + Map.of("country", "NL")); |
| 434 | + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 435 | + Map.of("country", "BG", "year", (short) 2023)); |
| 436 | + |
| 437 | + store.add(List.of(bgDocument, nlDocument, bgDocument2)); |
| 438 | + |
| 439 | + // Verify initial state |
| 440 | + List<Document> results = store |
| 441 | + .similaritySearch(SearchRequest.builder().query("The World").topK(5).build()); |
| 442 | + assertThat(results).hasSize(3); |
| 443 | + |
| 444 | + // Delete documents with country = BG |
| 445 | + Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ, |
| 446 | + new Filter.Key("country"), new Filter.Value("BG")); |
| 447 | + |
| 448 | + store.delete(filterExpression); |
| 449 | + |
| 450 | + results = store.similaritySearch( |
| 451 | + SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); |
| 452 | + |
| 453 | + assertThat(results).hasSize(1); |
| 454 | + assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); |
| 455 | + } |
| 456 | + }); |
| 457 | + } |
| 458 | + |
| 459 | + @Test |
| 460 | + void deleteWithStringFilterExpression() { |
| 461 | + this.contextRunner.run(context -> { |
| 462 | + try (CassandraVectorStore store = createTestStore(context, |
| 463 | + new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED), |
| 464 | + new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) { |
| 465 | + |
| 466 | + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 467 | + Map.of("country", "BG", "year", (short) 2020)); |
| 468 | + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 469 | + Map.of("country", "NL")); |
| 470 | + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", |
| 471 | + Map.of("country", "BG", "year", (short) 2023)); |
| 472 | + |
| 473 | + store.add(List.of(bgDocument, nlDocument, bgDocument2)); |
| 474 | + |
| 475 | + // Verify initial state |
| 476 | + List<Document> results = store |
| 477 | + .similaritySearch(SearchRequest.builder().query("The World").topK(5).build()); |
| 478 | + assertThat(results).hasSize(3); |
| 479 | + |
| 480 | + store.delete("country == 'BG'"); |
| 481 | + |
| 482 | + results = store.similaritySearch( |
| 483 | + SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build()); |
| 484 | + |
| 485 | + assertThat(results).hasSize(1); |
| 486 | + assertThat(results.get(0).getMetadata()).containsEntry("country", "NL"); |
| 487 | + } |
| 488 | + }); |
| 489 | + } |
| 490 | + |
| 491 | + @Test |
| 492 | + void deleteWithComplexFilterExpression() { |
| 493 | + this.contextRunner.run(context -> { |
| 494 | + try (CassandraVectorStore store = createTestStore(context, |
| 495 | + new SchemaColumn("type", DataTypes.TEXT, SchemaColumnTags.INDEXED), |
| 496 | + new SchemaColumn("priority", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) { |
| 497 | + |
| 498 | + var doc1 = new Document("Content 1", Map.of("type", "A", "priority", (short) 1)); |
| 499 | + var doc2 = new Document("Content 2", Map.of("type", "A", "priority", (short) 2)); |
| 500 | + var doc3 = new Document("Content 3", Map.of("type", "B", "priority", (short) 1)); |
| 501 | + |
| 502 | + store.add(List.of(doc1, doc2, doc3)); |
| 503 | + |
| 504 | + // Complex filter expression: (type == 'A' AND priority > 1) |
| 505 | + Filter.Expression priorityFilter = new Filter.Expression(Filter.ExpressionType.GT, |
| 506 | + new Filter.Key("priority"), new Filter.Value((short) 1)); |
| 507 | + Filter.Expression typeFilter = new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("type"), |
| 508 | + new Filter.Value("A")); |
| 509 | + Filter.Expression complexFilter = new Filter.Expression(Filter.ExpressionType.AND, typeFilter, |
| 510 | + priorityFilter); |
| 511 | + |
| 512 | + store.delete(complexFilter); |
| 513 | + |
| 514 | + var results = store.similaritySearch( |
| 515 | + SearchRequest.builder().query("Content").topK(5).similarityThresholdAll().build()); |
| 516 | + |
| 517 | + assertThat(results).hasSize(2); |
| 518 | + assertThat(results.stream().map(doc -> doc.getMetadata().get("type")).collect(Collectors.toList())) |
| 519 | + .containsExactlyInAnyOrder("A", "B"); |
| 520 | + assertThat(results.stream() |
| 521 | + .map(doc -> ((Short) doc.getMetadata().get("priority")).intValue()) |
| 522 | + .collect(Collectors.toList())).containsExactlyInAnyOrder(1, 1); |
| 523 | + } |
| 524 | + }); |
| 525 | + } |
| 526 | + |
420 | 527 | @SpringBootConfiguration
|
421 | 528 | @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
|
422 | 529 | public static class TestApplication {
|
|
0 commit comments