Request-time filter expressions for RAG

When using the RetrievalAugmentationAdvisor with the VectorStoreDocumentRetriever, it’s now possible to provide a filter expression at request-time as an advisor context variable with key VectorStoreDocumentRetriever.FILTER_EXPRESSION.

Fixes gh-1776

Signed-off-by: Thomas Vitale <[email protected]>

Author: ThomasVitale

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java

@@ -106,6 +106,7 @@ public AdvisedRequest before(AdvisedRequest request) {
 		Query originalQuery = Query.builder()
 			.text(new PromptTemplate(request.userText(), request.userParams()).render())
 			.history(request.messages())
+			.context(context)
 			.build();
 
 		// 1. Transform original user query based on a chain of query transformers.

spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java

@@ -24,8 +24,10 @@
 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 @@
  */
 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 VectorStoreDocumentRetriever(VectorStore vectorStore, @Nullable Double si
 	@Override
 	public List<Document> 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();
 	}

spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java

@@ -210,6 +210,30 @@ void retrieveWithQueryObjectAndDefaultValues() {
 		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<String> tenantIdentifier = new ThreadLocal<>();

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<Document> 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<Document> retrievedDocuments = documentRetriever.retrieve(query);
+----
+
 ==== Document Join
 
 A component for combining documents retrieved based on multiple queries and from multiple data sources into

spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java

@@ -108,6 +108,29 @@ void ragBasic() {
 		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();

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 @@ void whenTransformerWithDefaults() {
 
 		assertThat(transformedQuery).isNotNull();
 		System.out.println(transformedQuery);
-		assertThat(transformedQuery.text()).containsIgnoringCase("model");
+		assertThat(transformedQuery.text()).containsIgnoringCase("Angeles");
 	}
 
 }

spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java

@@ -46,20 +46,21 @@
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
 class VectorStoreDocumentRetrieverIT {
 
-	private static final Map<String, Document> 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<String, Document> 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 @@ void tearDown() {
 	}
 
 	@Test
-	void withFilter() {
+	void withBuildFilter() {
 		DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
 			.vectorStore(this.pgVectorStore)
 			.similarityThreshold(0.50)
@@ -95,7 +96,7 @@ void withFilter() {
 	}
 
 	@Test
-	void withNoFilter() {
+	void withNoBuildFilter() {
 		DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
 			.vectorStore(this.pgVectorStore)
 			.similarityThreshold(0.50)
@@ -110,4 +111,27 @@ void withNoFilter() {
 		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<Document> 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()));
+	}
+
 }