Modular RAG: Retrieval with Vector Stores

* Establish new package for Modular RAG components.
* Add new Query API, representing a query in the context of a RAG flow.
* Define Retrieval package for the RAG building blocks handling the data retrieval operations.
* Relocate DocumentRetriever to Retrieval package and implement VectorStoreDocumentRetriever.
* Introduce RetrievalAugmentationAdvisor as the successor of QuestionAnswerAdvisor. It uses the Retrieval building blocks described in the previous point.
* Make Advisor APIs null-safe and update tests accordingly.

Relates to gh-#1603

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>

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

@@ -0,0 +1,224 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.chat.client.advisor;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+
+import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
+import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
+import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisor;
+import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisorChain;
+import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisor;
+import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.PromptTemplate;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.model.Content;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.retrieval.source.DocumentRetriever;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
+
+/**
+ * This advisor implements common Retrieval Augmented Generation (RAG) flows using the
+ * building blocks defined in the {@link org.springframework.ai.rag} package and following
+ * the Modular RAG Architecture.
+ * <p>
+ * It's the successor of the {@link QuestionAnswerAdvisor}.
+ *
+ * @author Christian Tzolov
+ * @author Thomas Vitale
+ * @since 1.0.0
+ * @see <a href="http://export.arxiv.org/abs/2407.21059">arXiv:2407.21059</a>
+ * @see <a href="https://export.arxiv.org/abs/2312.10997">arXiv:2312.10997</a>
+ */
+public class RetrievalAugmentationAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {
+
+	public static final String DOCUMENT_CONTEXT = "rag_document_context";
+
+	public static final PromptTemplate DEFAULT_PROMPT_TEMPLATE = new PromptTemplate("""
+			{query}
+
+			Context information is below. Use this information to answer the user query.
+
+			---------------------
+			{context}
+			---------------------
+
+			Given the context and provided history information and not prior knowledge,
+			reply to the user query. If the answer is not in the context, inform
+			the user that you can't answer the query.
+			""");
+
+	private final DocumentRetriever documentRetriever;
+
+	private final PromptTemplate promptTemplate;
+
+	private final boolean protectFromBlocking;
+
+	private final int order;
+
+	public RetrievalAugmentationAdvisor(DocumentRetriever documentRetriever, @Nullable PromptTemplate promptTemplate,
+			@Nullable Boolean protectFromBlocking, @Nullable Integer order) {
+		Assert.notNull(documentRetriever, "documentRetriever cannot be null");
+		this.documentRetriever = documentRetriever;
+		this.promptTemplate = promptTemplate != null ? promptTemplate : DEFAULT_PROMPT_TEMPLATE;
+		this.protectFromBlocking = protectFromBlocking != null ? protectFromBlocking : false;
+		this.order = order != null ? order : 0;
+	}
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	@Override
+	public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
+		Assert.notNull(advisedRequest, "advisedRequest cannot be null");
+		Assert.notNull(chain, "chain cannot be null");
+
+		AdvisedRequest processedAdvisedRequest = before(advisedRequest);
+		AdvisedResponse advisedResponse = chain.nextAroundCall(processedAdvisedRequest);
+		return after(advisedResponse);
+	}
+
+	@Override
+	public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
+		Assert.notNull(advisedRequest, "advisedRequest cannot be null");
+		Assert.notNull(chain, "chain cannot be null");
+
+		// This can be executed by both blocking and non-blocking Threads
+		// E.g. a command line or Tomcat blocking Thread implementation
+		// or by a WebFlux dispatch in a non-blocking manner.
+		Flux<AdvisedResponse> advisedResponses = (this.protectFromBlocking) ?
+		// @formatter:off
+				Mono.just(advisedRequest)
+						.publishOn(Schedulers.boundedElastic())
+						.map(this::before)
+						.flatMapMany(chain::nextAroundStream)
+				: chain.nextAroundStream(before(advisedRequest));
+		// @formatter:on
+
+		return advisedResponses.map(ar -> {
+			if (onFinishReason().test(ar)) {
+				ar = after(ar);
+			}
+			return ar;
+		});
+	}
+
+	private AdvisedRequest before(AdvisedRequest request) {
+		Map<String, Object> context = new HashMap<>(request.adviseContext());
+
+		// 0. Create a query from the user text and parameters.
+		Query query = new Query(new PromptTemplate(request.userText(), request.userParams()).render());
+
+		// 1. Retrieve similar documents for the original query.
+		List<Document> documents = this.documentRetriever.retrieve(query);
+		context.put(DOCUMENT_CONTEXT, documents);
+
+		// 2. Combine retrieved documents.
+		String documentContext = documents.stream()
+			.map(Content::getContent)
+			.collect(Collectors.joining(System.lineSeparator()));
+
+		// 3. Define augmentation prompt parameters.
+		Map<String, Object> promptParameters = Map.of("query", query.text(), "context", documentContext);
+
+		// 4. Augment user prompt with the context data.
+		UserMessage augmentedUserMessage = (UserMessage) this.promptTemplate.createMessage(promptParameters);
+
+		return AdvisedRequest.from(request)
+			.withUserText(augmentedUserMessage.getContent())
+			.withAdviseContext(context)
+			.build();
+	}
+
+	private AdvisedResponse after(AdvisedResponse advisedResponse) {
+		ChatResponse.Builder chatResponseBuilder = ChatResponse.builder().from(advisedResponse.response());
+		chatResponseBuilder.withMetadata(DOCUMENT_CONTEXT, advisedResponse.adviseContext().get(DOCUMENT_CONTEXT));
+		return new AdvisedResponse(chatResponseBuilder.build(), advisedResponse.adviseContext());
+	}
+
+	private Predicate<AdvisedResponse> onFinishReason() {
+		return advisedResponse -> advisedResponse.response()
+			.getResults()
+			.stream()
+			.anyMatch(result -> result != null && result.getMetadata() != null
+					&& StringUtils.hasText(result.getMetadata().getFinishReason()));
+	}
+
+	@Override
+	public String getName() {
+		return this.getClass().getSimpleName();
+	}
+
+	@Override
+	public int getOrder() {
+		return this.order;
+	}
+
+	public static final class Builder {
+
+		private DocumentRetriever documentRetriever;
+
+		private PromptTemplate promptTemplate;
+
+		private Boolean protectFromBlocking;
+
+		private Integer order;
+
+		private Builder() {
+		}
+
+		public Builder documentRetriever(DocumentRetriever documentRetriever) {
+			this.documentRetriever = documentRetriever;
+			return this;
+		}
+
+		public Builder promptTemplate(PromptTemplate promptTemplate) {
+			this.promptTemplate = promptTemplate;
+			return this;
+		}
+
+		public Builder protectFromBlocking(Boolean protectFromBlocking) {
+			this.protectFromBlocking = protectFromBlocking;
+			return this;
+		}
+
+		public Builder order(Integer order) {
+			this.order = order;
+			return this;
+		}
+
+		public RetrievalAugmentationAdvisor build() {
+			return new RetrievalAugmentationAdvisor(this.documentRetriever, this.promptTemplate,
+					this.protectFromBlocking, this.order);
+		}
+
+	}
+
+}

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

@@ -34,6 +34,8 @@ import org.springframework.ai.chat.prompt.PromptTemplate;
 import org.springframework.ai.model.Media;
 import org.springframework.ai.model.function.FunctionCallback;
 import org.springframework.ai.model.function.FunctionCallingOptions;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
 
@@ -41,8 +43,6 @@ import org.springframework.util.StringUtils;
  * The data of the chat client request that can be modified before the execution of the
  * ChatClient's call method
  *
- * @author Christian Tzolov
- * @since 1.0.0
  * @param chatModel the chat model used
  * @param userText the text provided by the user
  * @param systemText the text provided by the system
@@ -57,13 +57,53 @@ import org.springframework.util.StringUtils;
  * @param advisorParams the map of advisor parameters
  * @param adviseContext the map of advise context
  * @param toolContext the tool context
+ * @author Christian Tzolov
+ * @author Thomas Vitale
+ * @since 1.0.0
  */
-public record AdvisedRequest(ChatModel chatModel, String userText, String systemText, ChatOptions chatOptions,
-		List<Media> media, List<String> functionNames, List<FunctionCallback> functionCallbacks, List<Message> messages,
-		Map<String, Object> userParams, Map<String, Object> systemParams, List<Advisor> advisors,
-		Map<String, Object> advisorParams, Map<String, Object> adviseContext, Map<String, Object> toolContext) {
+public record AdvisedRequest(
+// @formatter:off
+		ChatModel chatModel,
+		String userText,
+		@Nullable
+		String systemText,
+		@Nullable
+		ChatOptions chatOptions,
+		List<Media> media,
+		List<String> functionNames,
+		List<FunctionCallback> functionCallbacks,
+		List<Message> messages,
+		Map<String, Object> userParams,
+		Map<String, Object> systemParams,
+		List<Advisor> advisors,
+		Map<String, Object> advisorParams,
+		Map<String, Object> adviseContext,
+		Map<String, Object> toolContext
+// @formatter:on
+) {
+
+	public AdvisedRequest {
+		Assert.notNull(chatModel, "chatModel cannot be null");
+		Assert.hasText(userText, "userText cannot be null or empty");
+		Assert.notNull(media, "media cannot be null");
+		Assert.notNull(functionNames, "functionNames cannot be null");
+		Assert.notNull(functionCallbacks, "functionCallbacks cannot be null");
+		Assert.notNull(messages, "messages cannot be null");
+		Assert.notNull(userParams, "userParams cannot be null");
+		Assert.notNull(systemParams, "systemParams cannot be null");
+		Assert.notNull(advisors, "advisors cannot be null");
+		Assert.notNull(advisorParams, "advisorParams cannot be null");
+		Assert.notNull(adviseContext, "adviseContext cannot be null");
+		Assert.notNull(toolContext, "toolContext cannot be null");
+	}
+
+	public static Builder builder() {
+		return new Builder();
+	}
 
 	public static Builder from(AdvisedRequest from) {
+		Assert.notNull(from, "AdvisedRequest cannot be null");
+
 		Builder builder = new Builder();
 		builder.chatModel = from.chatModel;
 		builder.userText = from.userText;
@@ -79,23 +119,18 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
 		builder.advisorParams = from.advisorParams;
 		builder.adviseContext = from.adviseContext;
 		builder.toolContext = from.toolContext;
-
 		return builder;
 	}
 
-	public static Builder builder() {
-		return new Builder();
-	}
-
 	public AdvisedRequest updateContext(Function<Map<String, Object>, Map<String, Object>> contextTransform) {
+		Assert.notNull(contextTransform, "contextTransform cannot be null");
 		return from(this)
 			.withAdviseContext(Collections.unmodifiableMap(contextTransform.apply(new HashMap<>(this.adviseContext))))
 			.build();
 	}
 
 	public Prompt toPrompt() {
-
-		var messages = new ArrayList<Message>(this.messages());
+		var messages = new ArrayList<>(this.messages());
 
 		String processedSystemText = this.systemText();
 		if (StringUtils.hasText(processedSystemText)) {
@@ -111,7 +146,6 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
 				? this.userText() + System.lineSeparator() + "{spring_ai_soc_format}" : this.userText();
 
 		if (StringUtils.hasText(processedUserText)) {
-
 			Map<String, Object> userParams = new HashMap<>(this.userParams());
 			if (StringUtils.hasText(formatParam)) {
 				userParams.put("spring_ai_soc_format", formatParam);
@@ -137,17 +171,15 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
 		return new Prompt(messages, this.chatOptions());
 	}
 
-	public static class Builder {
-
-		public Map<String, Object> toolContext = Map.of();
+	public static final class Builder {
 
 		private ChatModel chatModel;
 
-		private String userText = "";
+		private String userText;
 
-		private String systemText = "";
+		private String systemText;
 
-		private ChatOptions chatOptions = null;
+		private ChatOptions chatOptions;
 
 		private List<Media> media = List.of();
 
@@ -167,6 +199,11 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
 
 		private Map<String, Object> adviseContext = Map.of();
 
+		public Map<String, Object> toolContext = Map.of();
+
+		private Builder() {
+		}
+
 		public Builder withChatModel(ChatModel chatModel) {
 			this.chatModel = chatModel;
 			return this;
@@ -202,11 +239,6 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
 			return this;
 		}
 
-		public Builder withToolContext(Map<String, Object> toolContext) {
-			this.toolContext = toolContext;
-			return this;
-		}
-
 		public Builder withMessages(List<Message> messages) {
 			this.messages = messages;
 			return this;
@@ -237,6 +269,11 @@ public record AdvisedRequest(ChatModel chatModel, String userText, String system
 			return this;
 		}
 
+		public Builder withToolContext(Map<String, Object> toolContext) {
+			this.toolContext = toolContext;
+			return this;
+		}
+
 		public AdvisedRequest build() {
 			return new AdvisedRequest(this.chatModel, this.userText, this.systemText, this.chatOptions, this.media,
 					this.functionNames, this.functionCallbacks, this.messages, this.userParams, this.systemParams,

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

@@ -25,48 +25,54 @@ import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.util.Assert;
 
 /**
+ * The data of the chat client response that can be modified before the call returns.
+ *
  * @author Christian Tzolov
+ * @author Thomas Vitale
  * @since 1.0.0
  */
 public record AdvisedResponse(ChatResponse response, Map<String, Object> adviseContext) {
 
+	public AdvisedResponse {
+		Assert.notNull(response, "response cannot be null");
+		Assert.notNull(adviseContext, "adviseContext cannot be null");
+	}
+
 	public static Builder builder() {
 		return new Builder();
 	}
 
+	public static Builder from(AdvisedResponse advisedResponse) {
+		Assert.notNull(advisedResponse, "advisedResponse cannot be null");
+		return new Builder().withResponse(advisedResponse.response).withAdviseContext(advisedResponse.adviseContext);
+	}
+
 	public AdvisedResponse updateContext(Function<Map<String, Object>, Map<String, Object>> contextTransform) {
+		Assert.notNull(contextTransform, "contextTransform cannot be null");
 		return new AdvisedResponse(this.response,
 				Collections.unmodifiableMap(contextTransform.apply(new HashMap<>(this.adviseContext))));
 	}
 
-	public static class Builder {
+	public static final class Builder {
 
 		private ChatResponse response;
 
 		private Map<String, Object> adviseContext;
 
-		public Builder() {
-		}
-
-		public static Builder from(AdvisedResponse advisedResponse) {
-			return new Builder().withResponse(advisedResponse.response)
-				.withAdviseContext(advisedResponse.adviseContext);
+		private Builder() {
 		}
 
 		public Builder withResponse(ChatResponse response) {
-			Assert.notNull(response, "the response must be non-null");
 			this.response = response;
 			return this;
 		}
 
 		public Builder withAdviseContext(Map<String, Object> adviseContext) {
-			Assert.notNull(adviseContext, "the adviseContext must be non-null");
 			this.adviseContext = adviseContext;
 			return this;
 		}
 
 		public AdvisedResponse build() {
-			Assert.notNull(this.adviseContext, "the adviseContext must be non-null");
 			return new AdvisedResponse(this.response, this.adviseContext);
 		}
 

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/package-info.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.chat.client.advisor.api;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.chat.client.advisor;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/Query.java

@@ -14,22 +14,18 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.document;
+package org.springframework.ai.rag;
 
-import java.util.List;
-import java.util.function.Function;
+import org.springframework.util.Assert;
 
-public interface DocumentRetriever extends Function<String, List<Document>> {
-
-	/**
-	 * Retrieves relevant documents however the implementation sees fit.
-	 * @param query query string
-	 * @return relevant documents
-	 */
-	List<Document> retrieve(String query);
-
-	default List<Document> apply(String query) {
-		return retrieve(query);
+/**
+ * Represents a query in the context of a Retrieval Augmented Generation (RAG) flow.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public record Query(String text) {
+	public Query {
+		Assert.hasText(text, "text cannot be null or empty");
 	}
-
 }

spring-ai-core/src/main/java/org/springframework/ai/rag/package-info.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * This package contains the core interfaces and classes supporting Retrieval Augmented
+ * Generation.
+ * <p>
+ * It's based on the Modular RAG Architecture and provides the necessary building blocks
+ * to define and execute RAG flows. It includes three levels of abstraction:
+ * <ol>
+ * <li>Module</li>
+ * <li>Sub-Module</li>
+ * <li>Operator</li>
+ * </ol>
+ *
+ * @see <a href="http://export.arxiv.org/abs/2407.21059">arXiv:2407.21059</a>
+ * @see <a href="https://export.arxiv.org/abs/2312.10997">arXiv:2312.10997</a>
+ */
+
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/package-info.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * RAG Module: Retrieval.
+ * <p>
+ * This package includes submodules for handling the retrieval process in RAG flows.
+ */
+
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.retrieval;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/source/DocumentRetriever.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.rag.retrieval.source;
+
+import java.util.List;
+import java.util.function.Function;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+
+/**
+ * API for retrieving {@link Document}s from an underlying data source.
+ *
+ * @author Christian Tzolov
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public interface DocumentRetriever extends Function<Query, List<Document>> {
+
+	/**
+	 * Retrieves {@link Document}s from an underlying data source using the given
+	 * {@link Query}.
+	 */
+	List<Document> retrieve(Query query);
+
+	/**
+	 * Retrieves {@link Document}s from an underlying data source using the given query
+	 * string.
+	 */
+	default List<Document> retrieve(String query) {
+		return retrieve(new Query(query));
+	}
+
+	/**
+	 * Retrieves {@link Document}s from an underlying data source using the given
+	 * {@link Query}.
+	 */
+	default List<Document> apply(Query query) {
+		return retrieve(query);
+	}
+
+}

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

@@ -0,0 +1,138 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.rag.retrieval.source;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+import org.springframework.ai.document.Document;
+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.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * A document retriever that uses a vector store to search for documents. It supports
+ * filtering based on metadata, similarity threshold, and top-k results.
+ *
+ * <p>
+ * Example usage: <pre>{@code
+ * VectorStoreDocumentRetriever retriever = VectorStoreDocumentRetriever.builder()
+ *     .vectorStore(vectorStore)
+ *     .similarityThreshold(0.73)
+ *     .topK(5)
+ *     .filterExpression(filterExpression)
+ *     .build();
+ * List<Document> documents = retriever.retrieve("example query");
+ * }</pre>
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ * @see VectorStore
+ * @see Filter.Expression
+ */
+public class VectorStoreDocumentRetriever implements DocumentRetriever {
+
+	private final VectorStore vectorStore;
+
+	private final Double similarityThreshold;
+
+	private final Integer topK;
+
+	// Supplier to allow for lazy evaluation of the filter expression,
+	// which may depend on the execution content. For example, you may want to
+	// filter dynamically based on the current user's identity or tenant ID.
+	private final Supplier<Filter.Expression> filterExpression;
+
+	public VectorStoreDocumentRetriever(VectorStore vectorStore, @Nullable Double similarityThreshold,
+			@Nullable Integer topK, @Nullable Supplier<Filter.Expression> filterExpression) {
+		Assert.notNull(vectorStore, "vectorStore cannot be null");
+		this.vectorStore = vectorStore;
+		this.similarityThreshold = similarityThreshold != null ? similarityThreshold
+				: SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL;
+		this.topK = topK != null ? topK : SearchRequest.DEFAULT_TOP_K;
+		this.filterExpression = filterExpression != null ? filterExpression : () -> null;
+	}
+
+	@Override
+	public List<Document> retrieve(Query query) {
+		Assert.notNull(query, "query cannot be null");
+		var searchRequest = SearchRequest.query(query.text())
+			.withFilterExpression(this.filterExpression.get())
+			.withSimilarityThreshold(this.similarityThreshold)
+			.withTopK(this.topK);
+		return this.vectorStore.similaritySearch(searchRequest);
+	}
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	/**
+	 * Builder for {@link VectorStoreDocumentRetriever}.
+	 */
+	public static final class Builder {
+
+		private VectorStore vectorStore;
+
+		private Double similarityThreshold;
+
+		private Integer topK;
+
+		private Supplier<Filter.Expression> filterExpression;
+
+		private Builder() {
+		}
+
+		public Builder vectorStore(VectorStore vectorStore) {
+			this.vectorStore = vectorStore;
+			return this;
+		}
+
+		public Builder similarityThreshold(Double similarityThreshold) {
+			Assert.notNull(similarityThreshold, "similarityThreshold cannot be null");
+			this.similarityThreshold = similarityThreshold;
+			return this;
+		}
+
+		public Builder topK(Integer topK) {
+			Assert.notNull(topK, "topK cannot be null");
+			Assert.isTrue(topK > 0, "topK must be greater than 0");
+			this.topK = topK;
+			return this;
+		}
+
+		public Builder filterExpression(Filter.Expression filterExpression) {
+			this.filterExpression = () -> filterExpression;
+			return this;
+		}
+
+		public Builder filterExpression(Supplier<Filter.Expression> filterExpression) {
+			this.filterExpression = filterExpression;
+			return this;
+		}
+
+		public VectorStoreDocumentRetriever build() {
+			return new VectorStoreDocumentRetriever(this.vectorStore, this.similarityThreshold, this.topK,
+					this.filterExpression);
+		}
+
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/source/package-info.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * RAG Sub-Module: Source.
+ * <p>
+ * This package provides the functional building blocks for retrieving documents from a
+ * data source.
+ */
+
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.retrieval.source;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/test/java/org/springframework/ai/chat/client/ChatClientTest.java

@@ -50,17 +50,12 @@ import static org.mockito.BDDMockito.given;
 
 /**
  * @author Christian Tzolov
+ * @author Thomas Vitale
  */
 @ExtendWith(MockitoExtension.class)
 public class ChatClientTest {
 
-	static Function<String, String> mockFunction = new Function<String, String>() {
-
-		@Override
-		public String apply(String s) {
-			return s;
-		}
-	};
+	static Function<String, String> mockFunction = s -> s;
 
 	@Mock
 	ChatModel chatModel;
@@ -88,7 +83,7 @@ public class ChatClientTest {
 
 		var chatClient = ChatClient.builder(this.chatModel).defaultSystem("Default system text").build();
 
-		var content = chatClient.prompt().call().content();
+		var content = chatClient.prompt("What's Spring AI?").call().content();
 
 		assertThat(content).isEqualTo("response");
 
@@ -96,7 +91,7 @@ public class ChatClientTest {
 		assertThat(systemMessage.getContent()).isEqualTo("Default system text");
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
-		content = join(chatClient.prompt().stream().content());
+		content = join(chatClient.prompt("What's Spring AI?").stream().content());
 
 		assertThat(content).isEqualTo("response");
 
@@ -105,7 +100,7 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// Override the default system text with prompt system
-		content = chatClient.prompt().system("Override default system text").call().content();
+		content = chatClient.prompt("What's Spring AI?").system("Override default system text").call().content();
 
 		assertThat(content).isEqualTo("response");
 		systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
@@ -113,7 +108,8 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// Streaming
-		content = join(chatClient.prompt().system("Override default system text").stream().content());
+		content = join(
+				chatClient.prompt("What's Spring AI?").system("Override default system text").stream().content());
 
 		assertThat(content).isEqualTo("response");
 		systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
@@ -140,7 +136,7 @@ public class ChatClientTest {
 				.param("param2", "value2"))
 			.build();
 
-		var content = chatClient.prompt().call().content();
+		var content = chatClient.prompt("What's Spring AI?").call().content();
 
 		assertThat(content).isEqualTo("response");
 
@@ -149,7 +145,7 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// Streaming
-		content = join(chatClient.prompt().stream().content());
+		content = join(chatClient.prompt("What's Spring AI?").stream().content());
 
 		assertThat(content).isEqualTo("response");
 
@@ -158,7 +154,7 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// Override single default system parameter
-		content = chatClient.prompt().system(s -> s.param("param1", "value1New")).call().content();
+		content = chatClient.prompt("What's Spring AI?").system(s -> s.param("param1", "value1New")).call().content();
 
 		assertThat(content).isEqualTo("response");
 		systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
@@ -166,7 +162,8 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// streaming
-		content = join(chatClient.prompt().system(s -> s.param("param1", "value1New")).stream().content());
+		content = join(
+				chatClient.prompt("What's Spring AI?").system(s -> s.param("param1", "value1New")).stream().content());
 
 		assertThat(content).isEqualTo("response");
 		systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
@@ -174,7 +171,7 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// Override default system text
-		content = chatClient.prompt()
+		content = chatClient.prompt("What's Spring AI?")
 			.system(s -> s.text("Override default system text {param3}").param("param3", "value3"))
 			.call()
 			.content();
@@ -185,7 +182,7 @@ public class ChatClientTest {
 		assertThat(systemMessage.getMessageType()).isEqualTo(MessageType.SYSTEM);
 
 		// Streaming
-		content = join(chatClient.prompt()
+		content = join(chatClient.prompt("What's Spring AI?")
 			.system(s -> s.text("Override default system text {param3}").param("param3", "value3"))
 			.stream()
 			.content());
@@ -489,11 +486,16 @@ public class ChatClientTest {
 		given(this.chatModel.call(this.promptCaptor.capture()))
 			.willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("response")))));
 
-		String response = ChatClient.builder(this.chatModel).build().prompt().system("System prompt").call().content();
+		String response = ChatClient.builder(this.chatModel)
+			.build()
+			.prompt("What's Spring AI?")
+			.system("System prompt")
+			.call()
+			.content();
 
 		assertThat(response).isEqualTo("response");
 
-		assertThat(this.promptCaptor.getValue().getInstructions()).hasSize(1);
+		assertThat(this.promptCaptor.getValue().getInstructions()).hasSize(2);
 
 		Message systemMessage = this.promptCaptor.getValue().getInstructions().get(0);
 		assertThat(systemMessage.getContent()).isEqualTo("System prompt");

spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisorTests.java

@@ -0,0 +1,112 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.chat.client.advisor;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import org.springframework.ai.chat.client.ChatClient;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.retrieval.source.DocumentRetriever;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Unit tests for {@link RetrievalAugmentationAdvisor}.
+ *
+ * @author Thomas Vitale
+ */
+class RetrievalAugmentationAdvisorTests {
+
+	@Test
+	void whenDocumentRetrieverIsNullThenThrow() {
+		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder().documentRetriever(null).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentRetriever cannot be null");
+	}
+
+	@Test
+	void theOneWithTheDocumentRetriever() {
+		// Chat Model
+		var chatModel = mock(ChatModel.class);
+		var promptCaptor = ArgumentCaptor.forClass(Prompt.class);
+		given(chatModel.call(promptCaptor.capture())).willReturn(ChatResponse.builder()
+			.withGenerations(List.of(new Generation(new AssistantMessage("Felix Felicis"))))
+			.build());
+
+		// Document Retriever
+		var documentContext = List.of(Document.builder().withId("1").withContent("doc1").build(),
+				Document.builder().withId("2").withContent("doc2").build());
+		var documentRetriever = mock(DocumentRetriever.class);
+		var queryCaptor = ArgumentCaptor.forClass(Query.class);
+		given(documentRetriever.retrieve(queryCaptor.capture())).willReturn(documentContext);
+
+		// Advisor
+		var advisor = RetrievalAugmentationAdvisor.builder().documentRetriever(documentRetriever).build();
+
+		// Chat Client
+		var chatClient = ChatClient.builder(chatModel)
+			.defaultAdvisors(advisor)
+			.defaultSystem("You are a wizard!")
+			.build();
+
+		// Call
+		var chatResponse = chatClient.prompt()
+			.user(user -> user.text("What would I get if I added {ingredient1} to {ingredient2}?")
+				.param("ingredient1", "a pinch of Moonstone")
+				.param("ingredient2", "a dash of powdered Gold"))
+			.call()
+			.chatResponse();
+
+		// Verify
+		assertThat(chatResponse.getResult().getOutput().getContent()).isEqualTo("Felix Felicis");
+		assertThat(chatResponse.getMetadata().<List<Document>>get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT))
+			.containsAll(documentContext);
+
+		var query = queryCaptor.getValue();
+		assertThat(query.text())
+			.isEqualTo("What would I get if I added a pinch of Moonstone to a dash of powdered Gold?");
+
+		var prompt = promptCaptor.getValue();
+		assertThat(prompt.getContents()).contains("""
+				What would I get if I added a pinch of Moonstone to a dash of powdered Gold?
+
+				Context information is below. Use this information to answer the user query.
+
+				---------------------
+				doc1
+				doc2
+				---------------------
+
+				Given the context and provided history information and not prior knowledge,
+				reply to the user query. If the answer is not in the context, inform
+				the user that you can't answer the query.
+				""");
+	}
+
+}

spring-ai-core/src/test/java/org/springframework/ai/rag/QueryTests.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.rag;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Unit tests for {@link Query}.
+ *
+ * @author Thomas Vitale
+ */
+class QueryTests {
+
+	@Test
+	void whenTextIsNullThenThrow() {
+		assertThatThrownBy(() -> new Query(null)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("text cannot be null or empty");
+	}
+
+	@Test
+	void whenTextIsEmptyThenThrow() {
+		assertThatThrownBy(() -> new Query("")).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("text cannot be null or empty");
+	}
+
+}

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

@@ -0,0 +1,208 @@
+/*
+ * Copyright 2023-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.ai.rag.retrieval.source;
+
+import java.util.List;
+import java.util.Map;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.internal.verification.Times;
+
+import org.springframework.ai.document.Document;
+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.FilterExpressionBuilder;
+import org.springframework.util.Assert;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ;
+
+/**
+ * Unit tests for {@link VectorStoreDocumentRetriever}.
+ */
+class VectorStoreDocumentRetrieverTests {
+
+	@Test
+	void whenVectorStoreIsNullThenThrow() {
+		assertThatThrownBy(() -> VectorStoreDocumentRetriever.builder().vectorStore(null).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("vectorStore cannot be null");
+	}
+
+	@Test
+	void searchRequestParameters() {
+		var mockVectorStore = mock(VectorStore.class);
+		var documentRetriever = VectorStoreDocumentRetriever.builder()
+			.vectorStore(mockVectorStore)
+			.similarityThreshold(0.73)
+			.topK(5)
+			.filterExpression(new Filter.Expression(EQ, new Filter.Key("location"), new Filter.Value("Rivendell")))
+			.build();
+
+		documentRetriever.retrieve("query");
+
+		var searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class);
+		verify(mockVectorStore).similaritySearch(searchRequestCaptor.capture());
+
+		var searchRequest = searchRequestCaptor.getValue();
+		assertThat(searchRequest.getQuery()).isEqualTo("query");
+		assertThat(searchRequest.getSimilarityThreshold()).isEqualTo(0.73);
+		assertThat(searchRequest.getTopK()).isEqualTo(5);
+		assertThat(searchRequest.getFilterExpression())
+			.isEqualTo(new Filter.Expression(EQ, new Filter.Key("location"), new Filter.Value("Rivendell")));
+	}
+
+	@Test
+	void dynamicFilterExpressions() {
+		var mockVectorStore = mock(VectorStore.class);
+		var documentRetriever = VectorStoreDocumentRetriever.builder()
+			.vectorStore(mockVectorStore)
+			.filterExpression(
+					() -> new FilterExpressionBuilder().eq("tenantId", TenantContextHolder.getTenantIdentifier())
+						.build())
+			.build();
+
+		TenantContextHolder.setTenantIdentifier("tenant1");
+		documentRetriever.retrieve("query");
+		TenantContextHolder.clear();
+
+		TenantContextHolder.setTenantIdentifier("tenant2");
+		documentRetriever.retrieve("query");
+		TenantContextHolder.clear();
+
+		var searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class);
+
+		verify(mockVectorStore, new Times(2)).similaritySearch(searchRequestCaptor.capture());
+
+		var searchRequest1 = searchRequestCaptor.getAllValues().get(0);
+		assertThat(searchRequest1.getFilterExpression())
+			.isEqualTo(new Filter.Expression(EQ, new Filter.Key("tenantId"), new Filter.Value("tenant1")));
+
+		var searchRequest2 = searchRequestCaptor.getAllValues().get(1);
+		assertThat(searchRequest2.getFilterExpression())
+			.isEqualTo(new Filter.Expression(EQ, new Filter.Key("tenantId"), new Filter.Value("tenant2")));
+	}
+
+	@Test
+	void whenQueryObjectIsNullThenThrow() {
+		var mockVectorStore = mock(VectorStore.class);
+		var documentRetriever = VectorStoreDocumentRetriever.builder().vectorStore(mockVectorStore).build();
+
+		Query nullQuery = null;
+		assertThatThrownBy(() -> documentRetriever.retrieve(nullQuery)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("query cannot be null");
+	}
+
+	@Test
+	void defaultValuesAreAppliedWhenNotSpecified() {
+		var mockVectorStore = mock(VectorStore.class);
+		var documentRetriever = VectorStoreDocumentRetriever.builder().vectorStore(mockVectorStore).build();
+
+		documentRetriever.retrieve("test query");
+
+		var searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class);
+		verify(mockVectorStore).similaritySearch(searchRequestCaptor.capture());
+
+		var searchRequest = searchRequestCaptor.getValue();
+		assertThat(searchRequest.getSimilarityThreshold()).isEqualTo(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL);
+		assertThat(searchRequest.getTopK()).isEqualTo(SearchRequest.DEFAULT_TOP_K);
+		assertThat(searchRequest.getFilterExpression()).isNull();
+	}
+
+	@Test
+	void retrieveWithQueryObject() {
+		var mockVectorStore = mock(VectorStore.class);
+		var documentRetriever = VectorStoreDocumentRetriever.builder()
+			.vectorStore(mockVectorStore)
+			.similarityThreshold(0.85)
+			.topK(3)
+			.filterExpression(new Filter.Expression(EQ, new Filter.Key("category"), new Filter.Value("books")))
+			.build();
+
+		var query = new Query("test query");
+		documentRetriever.retrieve(query);
+
+		var searchRequestCaptor = ArgumentCaptor.forClass(SearchRequest.class);
+		verify(mockVectorStore).similaritySearch(searchRequestCaptor.capture());
+
+		var searchRequest = searchRequestCaptor.getValue();
+		assertThat(searchRequest.getQuery()).isEqualTo("test query");
+		assertThat(searchRequest.getSimilarityThreshold()).isEqualTo(0.85);
+		assertThat(searchRequest.getTopK()).isEqualTo(3);
+		assertThat(searchRequest.getFilterExpression())
+			.isEqualTo(new Filter.Expression(EQ, new Filter.Key("category"), new Filter.Value("books")));
+	}
+
+	@Test
+	void retrieveWithQueryObjectAndDefaultValues() {
+		var mockVectorStore = mock(VectorStore.class);
+		var documentRetriever = VectorStoreDocumentRetriever.builder().vectorStore(mockVectorStore).build();
+
+		// Setup mock to return some documents
+		List<Document> mockDocuments = List.of(new Document("content1", Map.of("id", "1")),
+				new Document("content2", Map.of("id", "2")));
+		when(mockVectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments);
+
+		var query = new Query("test query");
+		var result = 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()).isNull();
+
+		// Verify the returned documents
+		assertThat(result).hasSize(2).containsExactlyElementsOf(mockDocuments);
+	}
+
+	static final class TenantContextHolder {
+
+		private static final ThreadLocal<String> tenantIdentifier = new ThreadLocal<>();
+
+		private TenantContextHolder() {
+		}
+
+		public static void setTenantIdentifier(String tenant) {
+			Assert.hasText(tenant, "tenant cannot be null or empty");
+			tenantIdentifier.set(tenant);
+		}
+
+		public static String getTenantIdentifier() {
+			return tenantIdentifier.get();
+		}
+
+		public static void clear() {
+			tenantIdentifier.remove();
+		}
+
+	}
+
+}

src/checkstyle/checkstyle.xml

@@ -158,13 +158,13 @@
 					  value="Line has leading space characters; indentation should be performed with tabs only." />
 			<property name="ignoreComments" value="true" />
 		</module>
-		<module name="com.puppycrawl.tools.checkstyle.checks.regexp.RegexpSinglelineJavaCheck">
-			<property name="maximum" value="0"/>
-			<property name="format" value="org\.mockito\.Mockito\.(when|doThrow|doAnswer)" />
-			<property name="message"
-					  value="Please use BDDMockito imports." />
-			<property name="ignoreComments" value="true" />
-		</module>
+<!--		<module name="com.puppycrawl.tools.checkstyle.checks.regexp.RegexpSinglelineJavaCheck">-->
+<!--			<property name="maximum" value="0"/>-->
+<!--			<property name="format" value="org\.mockito\.Mockito\.(when|doThrow|doAnswer)" />-->
+<!--			<property name="message"-->
+<!--					  value="Please use BDDMockito imports." />-->
+<!--			<property name="ignoreComments" value="true" />-->
+<!--		</module>-->
 		<module name="com.puppycrawl.tools.checkstyle.checks.regexp.RegexpSinglelineJavaCheck">
 			<property name="maximum" value="0"/>
 			<property name="format" value="org\.junit\.Assert\.assert" />