diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java index 3db204aca..a76fc55c2 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java @@ -24,8 +24,10 @@ import org.springframework.ai.rag.Query; 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; import org.springframework.lang.Nullable; import org.springframework.util.Assert; +import org.springframework.util.StringUtils; /** * Retrieves documents from a vector store that are semantically similar to the input @@ -48,6 +50,8 @@ import org.springframework.util.Assert; */ public final class VectorStoreDocumentRetriever implements DocumentRetriever { + public static final String FILTER_EXPRESSION = "vector_store_filter_expression"; + private final VectorStore vectorStore; private final Double similarityThreshold; @@ -75,15 +79,24 @@ public final class VectorStoreDocumentRetriever implements DocumentRetriever { @Override public List retrieve(Query query) { Assert.notNull(query, "query cannot be null"); + var requestFilterExpression = computeRequestFilterExpression(query); var searchRequest = SearchRequest.builder() .query(query.text()) - .filterExpression(this.filterExpression.get()) + .filterExpression(requestFilterExpression) .similarityThreshold(this.similarityThreshold) .topK(this.topK) .build(); return this.vectorStore.similaritySearch(searchRequest); } + private Filter.Expression computeRequestFilterExpression(Query query) { + var contextFilterExpression = query.context().get(FILTER_EXPRESSION); + if (contextFilterExpression != null && StringUtils.hasText(contextFilterExpression.toString())) { + return new FilterExpressionTextParser().parse(contextFilterExpression.toString()); + } + return this.filterExpression.get(); + } + public static Builder builder() { return new Builder(); } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java index 9cb4ea043..bbcc7f656 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java @@ -210,6 +210,30 @@ class VectorStoreDocumentRetrieverTests { assertThat(result).hasSize(2).containsExactlyElementsOf(mockDocuments); } + @Test + void retrieveWithQueryObjectAndRequestFilterExpression() { + var mockVectorStore = mock(VectorStore.class); + var documentRetriever = VectorStoreDocumentRetriever.builder().vectorStore(mockVectorStore).build(); + + var query = Query.builder() + .text("test query") + .context(Map.of(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "location == 'Rivendell'")) + .build(); + documentRetriever.retrieve(query); + + // Verify the mock interaction + var searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class); + verify(mockVectorStore).similaritySearch(searchRequestCaptor.capture()); + + // Verify the search request + var searchRequest = searchRequestCaptor.getValue(); + assertThat(searchRequest.getQuery()).isEqualTo("test query"); + assertThat(searchRequest.getSimilarityThreshold()).isEqualTo(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL); + assertThat(searchRequest.getTopK()).isEqualTo(SearchRequest.DEFAULT_TOP_K); + assertThat(searchRequest.getFilterExpression()) + .isEqualTo(new FilterExpressionBuilder().eq("location", "Rivendell").build()); + } + static final class TenantContextHolder { private static final ThreadLocal tenantIdentifier = new ThreadLocal<>(); diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/retrieval-augmented-generation.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/retrieval-augmented-generation.adoc index 2cbc22118..44d108742 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/retrieval-augmented-generation.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/retrieval-augmented-generation.adoc @@ -39,15 +39,12 @@ This filter expression can be configured when creating the `QuestionAnswerAdviso Here is how to create an instance of `QuestionAnswerAdvisor` where the threshold is `0.8` and to return the top `6` reulsts. - [source,java] ---- var qaAdvisor = new QuestionAnswerAdvisor(this.vectorStore, SearchRequest.builder().similarityThreshold(0.8d).topK(6).build()); ---- - - ==== Dynamic Filter Expressions Update the `SearchRequest` filter expression at runtime using the `FILTER_EXPRESSION` advisor context parameter: @@ -118,6 +115,29 @@ String answer = chatClient.prompt() .content(); ---- +The `VectorStoreDocumentRetriever` accepts a `FilterExpression` to filter the search results based on metadata. +You can provide one when instantiating the `VectorStoreDocumentRetriever` or at runtime per request, +using the `FILTER_EXPRESSION` advisor context parameter. + +[source,java] +---- +Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder() + .documentRetriever(VectorStoreDocumentRetriever.builder() + .similarityThreshold(0.50) + .vectorStore(vectorStore) + .build()) + .build(); + +String answer = chatClient.prompt() + .advisors(retrievalAugmentationAdvisor) + .advisors(a -> a.param(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "type == 'Spring'")) + .user(question) + .call() + .content(); +---- + +See xref:api/retrieval-augmented-generation.adoc#_vectorstoredocumentretriever for more information. + ===== Advanced RAG [source,java] @@ -298,6 +318,18 @@ DocumentRetriever retriever = VectorStoreDocumentRetriever.builder() List documents = retriever.retrieve(new Query("What are the KPIs for the next semester?")); ---- +You can also provide a request-specific filter expression via the `Query` API, using the `FILTER_EXPRESSION` parameter. +If both the request-specific and the retriever-specific filter expressions are provided, the request-specific filter expression takes precedence. + +[source,java] +---- +Query query = Query.builder() + .text("Who is Anacletus?") + .context(Map.of(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "location == 'Whispering Woods'")) + .build(); +List retrievedDocuments = documentRetriever.retrieve(query); +---- + ==== Document Join A component for combining documents retrieved based on multiple queries and from multiple data sources into diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java index 7e4e7afdf..e7170c68b 100644 --- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java +++ b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java @@ -108,6 +108,29 @@ class RetrievalAugmentationAdvisorIT { evaluateRelevancy(question, chatResponse); } + @Test + void ragWithRequestFilter() { + String question = "Where does the adventure of Anacletus and Birba take place?"; + + RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder() + .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build()) + .build(); + + ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel) + .build() + .prompt(question) + .advisors(ragAdvisor) + .advisors(a -> a.param(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "location == 'Italy'")) + .call() + .chatResponse(); + + assertThat(chatResponse).isNotNull(); + // No documents retrieved since the filter expression matches none of the + // documents in the vector store. + assertThat((String) chatResponse.getResult().getMetadata().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT)) + .isNull(); + } + @Test void ragWithCompression() { MessageChatMemoryAdvisor memoryAdvisor = MessageChatMemoryAdvisor.builder(new InMemoryChatMemory()).build(); diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/RewriteQueryTransformerIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/RewriteQueryTransformerIT.java index 32c93fc65..61be322c4 100644 --- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/RewriteQueryTransformerIT.java +++ b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/RewriteQueryTransformerIT.java @@ -43,7 +43,7 @@ class RewriteQueryTransformerIT { @Test void whenTransformerWithDefaults() { - Query query = new Query("I'm studying machine learning. What is an LLM?"); + Query query = new Query("What are the main tourist attractions in L.A.?"); QueryTransformer queryTransformer = RewriteQueryTransformer.builder() .chatClientBuilder(ChatClient.builder(this.openAiChatModel)) .build(); @@ -52,7 +52,7 @@ class RewriteQueryTransformerIT { assertThat(transformedQuery).isNotNull(); System.out.println(transformedQuery); - assertThat(transformedQuery.text()).containsIgnoringCase("model"); + assertThat(transformedQuery.text()).containsIgnoringCase("Angeles"); } } diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java index 3a90627a6..00fb7d7da 100644 --- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java +++ b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java @@ -46,20 +46,21 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*") class VectorStoreDocumentRetrieverIT { - private static final Map documents = Map.of("1", new Document( - "Anacletus was a majestic snowy owl with unusually bright golden eyes and distinctive black speckles across his wings.", - Map.of("location", "Whispering Woods")), "2", - new Document( + // @formatter:off + private static final Map documents = Map.of( + "1", new Document( + "Anacletus was a majestic snowy owl with unusually bright golden eyes and distinctive black speckles across his wings.", + Map.of("location", "Whispering Woods")), + "2", new Document( "Anacletus made his home in an ancient hollow oak tree deep within the Whispering Woods, where local villagers often heard his haunting calls at midnight.", Map.of("location", "Whispering Woods")), - "3", - new Document( + "3", new Document( "Despite being a nocturnal hunter like other owls, Anacletus had developed a peculiar habit of collecting shiny objects, especially lost coins and jewelry that glinted in the moonlight.", Map.of()), - "4", - new Document( + "4", new Document( "Birba was a plump Siamese cat with mismatched eyes - one blue and one green - who spent her days lounging on velvet cushions and judging everyone with a perpetual look of disdain.", Map.of("location", "Alfea"))); + // @formatter:on @Autowired PgVectorStore pgVectorStore; @@ -75,7 +76,7 @@ class VectorStoreDocumentRetrieverIT { } @Test - void withFilter() { + void withBuildFilter() { DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder() .vectorStore(this.pgVectorStore) .similarityThreshold(0.50) @@ -95,7 +96,7 @@ class VectorStoreDocumentRetrieverIT { } @Test - void withNoFilter() { + void withNoBuildFilter() { DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder() .vectorStore(this.pgVectorStore) .similarityThreshold(0.50) @@ -110,4 +111,27 @@ class VectorStoreDocumentRetrieverIT { assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("3").getId())); } + @Test + void withRequestFilter() { + DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder() + .vectorStore(this.pgVectorStore) + .similarityThreshold(0.50) + .topK(3) + .build(); + + Query query = Query.builder() + .text("Who is Anacletus?") + .context(Map.of(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "location == 'Whispering Woods'")) + .build(); + List retrievedDocuments = documentRetriever.retrieve(query); + + assertThat(retrievedDocuments).hasSize(2); + assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("1").getId())); + assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("2").getId())); + + // No request filter expression applied, so full access to all documents. + retrievedDocuments = documentRetriever.retrieve(new Query("Who is Birba?")); + assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("4").getId())); + } + }