Modular RAG: Orchestration and Post-Retrieval

Pre-Retrieval:
* Consolidated naming and documentation

Retrieval:
* Consolidated naming and documentation
* Introduced DocumentJoiner sub-module and CompositionDocumentJoiner operator

Post-Retrieval:
* Introduced main interfaces for sub-modules. Implementation waiting for missing features in Document APIs

Orchestration:
* Introduced QueryRouter sub-module and AllDocumentRetrieversQueryRouter operator

Generation:
* Consolidated naming and documentation

Advisor:
* Introduced BaseAdvisor to reduce boilerplate when implementing Advisors
* Extended RetrievalAugmentationAdvisor to include the new sub-modules

Relates to #gh-1603

spring-ai-core/pom.xml

@@ -96,6 +96,11 @@
 			<artifactId>micrometer-core</artifactId>
 		</dependency>
 
+		<dependency>
+			<groupId>io.micrometer</groupId>
+			<artifactId>context-propagation</artifactId>
+		</dependency>
+
 		<dependency>
 			<groupId>io.micrometer</groupId>
 			<artifactId>micrometer-tracing-bridge-otel</artifactId>
@@ -195,4 +200,4 @@
 	</profiles>
 
 
-</project>
+</project>

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

@@ -16,41 +16,41 @@
 
 package org.springframework.ai.chat.client.advisor;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Predicate;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
 
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-import reactor.core.scheduler.Schedulers;
+import reactor.core.scheduler.Scheduler;
 
 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.client.advisor.api.BaseAdvisor;
 import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.ai.chat.prompt.PromptTemplate;
 import org.springframework.ai.document.Document;
 import org.springframework.ai.rag.Query;
-import org.springframework.ai.rag.analysis.query.transformation.QueryTransformer;
-import org.springframework.ai.rag.augmentation.ContextualQueryAugmentor;
-import org.springframework.ai.rag.augmentation.QueryAugmentor;
+import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
+import org.springframework.ai.rag.generation.augmentation.QueryAugmenter;
+import org.springframework.ai.rag.orchestration.routing.AllRetrieversQueryRouter;
+import org.springframework.ai.rag.orchestration.routing.QueryRouter;
+import org.springframework.ai.rag.preretrieval.query.expansion.QueryExpander;
+import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
+import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
+import org.springframework.ai.rag.retrieval.join.DocumentJoiner;
 import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
+import org.springframework.core.task.TaskExecutor;
+import org.springframework.core.task.support.ContextPropagatingTaskDecorator;
 import org.springframework.lang.Nullable;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.util.Assert;
-import org.springframework.util.StringUtils;
 
 /**
  * Advisor that 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
@@ -58,29 +58,40 @@ import org.springframework.util.StringUtils;
  * @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 final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {
+public final class RetrievalAugmentationAdvisor implements BaseAdvisor {
 
 	public static final String DOCUMENT_CONTEXT = "rag_document_context";
 
 	private final List<QueryTransformer> queryTransformers;
 
-	private final DocumentRetriever documentRetriever;
+	@Nullable
+	private final QueryExpander queryExpander;
 
-	private final QueryAugmentor queryAugmentor;
+	private final QueryRouter queryRouter;
 
-	private final boolean protectFromBlocking;
+	private final DocumentJoiner documentJoiner;
+
+	private final QueryAugmenter queryAugmenter;
+
+	private final TaskExecutor taskExecutor;
+
+	private final Scheduler scheduler;
 
 	private final int order;
 
-	public RetrievalAugmentationAdvisor(List<QueryTransformer> queryTransformers, DocumentRetriever documentRetriever,
-			@Nullable QueryAugmentor queryAugmentor, @Nullable Boolean protectFromBlocking, @Nullable Integer order) {
-		Assert.notNull(queryTransformers, "queryTransformers cannot be null");
+	public RetrievalAugmentationAdvisor(@Nullable List<QueryTransformer> queryTransformers,
+			@Nullable QueryExpander queryExpander, QueryRouter queryRouter, @Nullable DocumentJoiner documentJoiner,
+			@Nullable QueryAugmenter queryAugmenter, @Nullable TaskExecutor taskExecutor, @Nullable Scheduler scheduler,
+			@Nullable Integer order) {
+		Assert.notNull(queryRouter, "queryRouter cannot be null");
 		Assert.noNullElements(queryTransformers, "queryTransformers cannot contain null elements");
-		Assert.notNull(documentRetriever, "documentRetriever cannot be null");
-		this.queryTransformers = queryTransformers;
-		this.documentRetriever = documentRetriever;
-		this.queryAugmentor = queryAugmentor != null ? queryAugmentor : ContextualQueryAugmentor.builder().build();
-		this.protectFromBlocking = protectFromBlocking != null ? protectFromBlocking : true;
+		this.queryTransformers = queryTransformers != null ? queryTransformers : List.of();
+		this.queryExpander = queryExpander;
+		this.queryRouter = queryRouter;
+		this.documentJoiner = documentJoiner != null ? documentJoiner : new ConcatenationDocumentJoiner();
+		this.queryAugmenter = queryAugmenter != null ? queryAugmenter : ContextualQueryAugmenter.builder().build();
+		this.taskExecutor = taskExecutor != null ? taskExecutor : buildDefaultTaskExecutor();
+		this.scheduler = scheduler != null ? scheduler : BaseAdvisor.DEFAULT_SCHEDULER;
 		this.order = order != null ? order : 0;
 	}
 
@@ -89,41 +100,7 @@ public final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, St
 	}
 
 	@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) {
+	public AdvisedRequest before(AdvisedRequest request) {
 		Map<String, Object> context = new HashMap<>(request.adviseContext());
 
 		// 0. Create a query from the user text and parameters.
@@ -135,17 +112,47 @@ public final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, St
 			transformedQuery = queryTransformer.apply(transformedQuery);
 		}
 
-		// 2. Retrieve similar documents for the original query.
-		List<Document> documents = this.documentRetriever.retrieve(transformedQuery);
+		// 2. Expand query into one or multiple queries.
+		List<Query> expandedQueries = this.queryExpander != null ? this.queryExpander.expand(transformedQuery)
+				: List.of(transformedQuery);
+
+		// 3. Get similar documents for each query.
+		Map<Query, List<List<Document>>> documentsForQuery = expandedQueries.stream()
+			.map(query -> CompletableFuture.supplyAsync(() -> getDocumentsForQuery(query), this.taskExecutor))
+			.toList()
+			.stream()
+			.map(CompletableFuture::join)
+			.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+		// 4. Combine documents retrieved based on multiple queries and from multiple data
+		// sources.
+		List<Document> documents = this.documentJoiner.join(documentsForQuery);
 		context.put(DOCUMENT_CONTEXT, documents);
 
-		// 3. Augment user query with the document contextual data.
-		Query augmentedQuery = this.queryAugmentor.augment(transformedQuery, documents);
+		// 5. Augment user query with the document contextual data.
+		Query augmentedQuery = this.queryAugmenter.augment(originalQuery, documents);
 
+		// 6. Update advised request with augmented prompt.
 		return AdvisedRequest.from(request).withUserText(augmentedQuery.text()).withAdviseContext(context).build();
 	}
 
-	private AdvisedResponse after(AdvisedResponse advisedResponse) {
+	/**
+	 * Processes a single query by routing it to document retrievers and collecting
+	 * documents.
+	 */
+	private Map.Entry<Query, List<List<Document>>> getDocumentsForQuery(Query query) {
+		List<DocumentRetriever> retrievers = this.queryRouter.route(query);
+		List<List<Document>> documents = retrievers.stream()
+			.map(retriever -> CompletableFuture.supplyAsync(() -> retriever.retrieve(query), this.taskExecutor))
+			.toList()
+			.stream()
+			.map(CompletableFuture::join)
+			.toList();
+		return Map.entry(query, documents);
+	}
+
+	@Override
+	public AdvisedResponse after(AdvisedResponse advisedResponse) {
 		ChatResponse.Builder chatResponseBuilder;
 		if (advisedResponse.response() == null) {
 			chatResponseBuilder = ChatResponse.builder();
@@ -157,20 +164,9 @@ public final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, St
 		return new AdvisedResponse(chatResponseBuilder.build(), advisedResponse.adviseContext());
 	}
 
-	private Predicate<AdvisedResponse> onFinishReason() {
-		return advisedResponse -> {
-			ChatResponse chatResponse = advisedResponse.response();
-			return chatResponse != null && chatResponse.getResults() != null
-					&& chatResponse.getResults()
-						.stream()
-						.anyMatch(result -> result != null && result.getMetadata() != null
-								&& StringUtils.hasText(result.getMetadata().getFinishReason()));
-		};
-	}
-
 	@Override
-	public String getName() {
-		return this.getClass().getSimpleName();
+	public Scheduler getScheduler() {
+		return this.scheduler;
 	}
 
 	@Override
@@ -178,15 +174,31 @@ public final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, St
 		return this.order;
 	}
 
+	private static TaskExecutor buildDefaultTaskExecutor() {
+		ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
+		taskExecutor.setThreadNamePrefix("ai-advisor-");
+		taskExecutor.setCorePoolSize(4);
+		taskExecutor.setMaxPoolSize(16);
+		taskExecutor.setTaskDecorator(new ContextPropagatingTaskDecorator());
+		taskExecutor.initialize();
+		return taskExecutor;
+	}
+
 	public static final class Builder {
 
-		private final List<QueryTransformer> queryTransformers = new ArrayList<>();
+		private List<QueryTransformer> queryTransformers;
 
-		private DocumentRetriever documentRetriever;
+		private QueryExpander queryExpander;
 
-		private QueryAugmentor queryAugmentor;
+		private QueryRouter queryRouter;
 
-		private Boolean protectFromBlocking;
+		private DocumentJoiner documentJoiner;
+
+		private QueryAugmenter queryAugmenter;
+
+		private TaskExecutor taskExecutor;
+
+		private Scheduler scheduler;
 
 		private Integer order;
 
@@ -194,29 +206,49 @@ public final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, St
 		}
 
 		public Builder queryTransformers(List<QueryTransformer> queryTransformers) {
-			Assert.notNull(queryTransformers, "queryTransformers cannot be null");
-			this.queryTransformers.addAll(queryTransformers);
+			this.queryTransformers = queryTransformers;
 			return this;
 		}
 
 		public Builder queryTransformers(QueryTransformer... queryTransformers) {
-			Assert.notNull(queryTransformers, "queryTransformers cannot be null");
-			this.queryTransformers.addAll(Arrays.asList(queryTransformers));
+			this.queryTransformers = Arrays.asList(queryTransformers);
+			return this;
+		}
+
+		public Builder queryExpander(QueryExpander queryExpander) {
+			this.queryExpander = queryExpander;
+			return this;
+		}
+
+		public Builder queryRouter(QueryRouter queryRouter) {
+			Assert.isNull(this.queryRouter, "Cannot set both documentRetriever and queryRouter");
+			this.queryRouter = queryRouter;
 			return this;
 		}
 
 		public Builder documentRetriever(DocumentRetriever documentRetriever) {
-			this.documentRetriever = documentRetriever;
+			Assert.isNull(this.queryRouter, "Cannot set both documentRetriever and queryRouter");
+			this.queryRouter = AllRetrieversQueryRouter.builder().documentRetrievers(documentRetriever).build();
 			return this;
 		}
 
-		public Builder queryAugmentor(QueryAugmentor queryAugmentor) {
-			this.queryAugmentor = queryAugmentor;
+		public Builder documentJoiner(DocumentJoiner documentJoiner) {
+			this.documentJoiner = documentJoiner;
 			return this;
 		}
 
-		public Builder protectFromBlocking(Boolean protectFromBlocking) {
-			this.protectFromBlocking = protectFromBlocking;
+		public Builder queryAugmenter(QueryAugmenter queryAugmenter) {
+			this.queryAugmenter = queryAugmenter;
+			return this;
+		}
+
+		public Builder taskExecutor(TaskExecutor taskExecutor) {
+			this.taskExecutor = taskExecutor;
+			return this;
+		}
+
+		public Builder scheduler(Scheduler scheduler) {
+			this.scheduler = scheduler;
 			return this;
 		}
 
@@ -226,8 +258,8 @@ public final class RetrievalAugmentationAdvisor implements CallAroundAdvisor, St
 		}
 
 		public RetrievalAugmentationAdvisor build() {
-			return new RetrievalAugmentationAdvisor(this.queryTransformers, this.documentRetriever, this.queryAugmentor,
-					this.protectFromBlocking, this.order);
+			return new RetrievalAugmentationAdvisor(this.queryTransformers, this.queryExpander, this.queryRouter,
+					this.documentJoiner, this.queryAugmenter, this.taskExecutor, this.scheduler, this.order);
 		}
 
 	}

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

@@ -0,0 +1,108 @@
+/*
+ * 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.api;
+
+import java.util.function.Predicate;
+
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
+
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
+
+/**
+ * Base advisor that implements common aspects of the {@link CallAroundAdvisor} and
+ * {@link StreamAroundAdvisor}, reducing the boilerplate code needed to implement an
+ * advisor. It provides default implementations for the
+ * {@link #aroundCall(AdvisedRequest, CallAroundAdvisorChain)} and
+ * {@link #aroundStream(AdvisedRequest, StreamAroundAdvisorChain)} methods, delegating the
+ * actual logic to the {@link #before(AdvisedRequest)} and {@link #after(AdvisedResponse)}
+ * methods.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public interface BaseAdvisor extends CallAroundAdvisor, StreamAroundAdvisor {
+
+	Scheduler DEFAULT_SCHEDULER = Schedulers.boundedElastic();
+
+	@Override
+	default 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
+	default Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
+		Assert.notNull(advisedRequest, "advisedRequest cannot be null");
+		Assert.notNull(chain, "chain cannot be null");
+		Assert.notNull(getScheduler(), "scheduler cannot be null");
+
+		Flux<AdvisedResponse> advisedResponses = Mono.just(advisedRequest)
+			.publishOn(getScheduler())
+			.map(this::before)
+			.flatMapMany(chain::nextAroundStream);
+
+		return advisedResponses.map(ar -> {
+			if (onFinishReason().test(ar)) {
+				ar = after(ar);
+			}
+			return ar;
+		}).onErrorResume(error -> Flux.error(new IllegalStateException("Stream processing failed", error)));
+	}
+
+	private Predicate<AdvisedResponse> onFinishReason() {
+		return advisedResponse -> {
+			ChatResponse chatResponse = advisedResponse.response();
+			return chatResponse != null && chatResponse.getResults() != null
+					&& chatResponse.getResults()
+						.stream()
+						.anyMatch(result -> result != null && result.getMetadata() != null
+								&& StringUtils.hasText(result.getMetadata().getFinishReason()));
+		};
+	}
+
+	@Override
+	default String getName() {
+		return this.getClass().getSimpleName();
+	}
+
+	/**
+	 * Logic to be executed before the rest of the advisor chain is called.
+	 */
+	AdvisedRequest before(AdvisedRequest request);
+
+	/**
+	 * Logic to be executed after the rest of the advisor chain is called.
+	 */
+	AdvisedResponse after(AdvisedResponse advisedResponse);
+
+	/**
+	 * Scheduler used for processing the advisor logic when streaming.
+	 */
+	default Scheduler getScheduler() {
+		return DEFAULT_SCHEDULER;
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/ContextualQueryAugmenter.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.augmentation;
+package org.springframework.ai.rag.generation.augmentation;
 
 import java.util.List;
 import java.util.Map;
@@ -37,18 +37,18 @@ import org.springframework.util.Assert;
  *
  * <p>
  * Example usage: <pre>{@code
- * QueryAugmentor augmentor = ContextualQueryAugmentor.builder()
+ * QueryAugmenter augmenter = ContextualQueryAugmenter.builder()
  *    .allowEmptyContext(false)
  *    .build();
- * Query augmentedQuery = augmentor.augment(query, documents);
+ * Query augmentedQuery = augmenter.augment(query, documents);
  * }</pre>
  *
  * @author Thomas Vitale
  * @since 1.0.0
  */
-public final class ContextualQueryAugmentor implements QueryAugmentor {
+public final class ContextualQueryAugmenter implements QueryAugmenter {
 
-	private static final Logger logger = LoggerFactory.getLogger(ContextualQueryAugmentor.class);
+	private static final Logger logger = LoggerFactory.getLogger(ContextualQueryAugmenter.class);
 
 	private static final PromptTemplate DEFAULT_PROMPT_TEMPLATE = new PromptTemplate("""
 			Context information is below.
@@ -74,7 +74,7 @@ public final class ContextualQueryAugmentor implements QueryAugmentor {
 			Politely inform the user that you can't answer it.
 			""");
 
-	private static final boolean DEFAULT_ALLOW_EMPTY_CONTEXT = true;
+	private static final boolean DEFAULT_ALLOW_EMPTY_CONTEXT = false;
 
 	private final PromptTemplate promptTemplate;
 
@@ -82,7 +82,7 @@ public final class ContextualQueryAugmentor implements QueryAugmentor {
 
 	private final boolean allowEmptyContext;
 
-	public ContextualQueryAugmentor(@Nullable PromptTemplate promptTemplate,
+	public ContextualQueryAugmenter(@Nullable PromptTemplate promptTemplate,
 			@Nullable PromptTemplate emptyContextPromptTemplate, @Nullable Boolean allowEmptyContext) {
 		this.promptTemplate = promptTemplate != null ? promptTemplate : DEFAULT_PROMPT_TEMPLATE;
 		this.emptyContextPromptTemplate = emptyContextPromptTemplate != null ? emptyContextPromptTemplate
@@ -102,7 +102,7 @@ public final class ContextualQueryAugmentor implements QueryAugmentor {
 			return augmentQueryWhenEmptyContext(query);
 		}
 
-		// 1. Join documents.
+		// 1. Collect content from documents.
 		String documentContext = documents.stream()
 			.map(Content::getContent)
 			.collect(Collectors.joining(System.lineSeparator()));
@@ -150,8 +150,8 @@ public final class ContextualQueryAugmentor implements QueryAugmentor {
 			return this;
 		}
 
-		public ContextualQueryAugmentor build() {
-			return new ContextualQueryAugmentor(this.promptTemplate, this.emptyContextPromptTemplate,
+		public ContextualQueryAugmenter build() {
+			return new ContextualQueryAugmenter(this.promptTemplate, this.emptyContextPromptTemplate,
 					this.allowEmptyContext);
 		}
 

spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/QueryAugmenter.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.augmentation;
+package org.springframework.ai.rag.generation.augmentation;
 
 import java.util.List;
 import java.util.function.BiFunction;
@@ -23,13 +23,13 @@ import org.springframework.ai.document.Document;
 import org.springframework.ai.rag.Query;
 
 /**
- * Component responsible for augmenting an input query with additional contextual data
- * that can be used by a large language model to answer the query.
+ * A component for augmenting an input query with additional data, useful to provide a
+ * large language model with the necessary context to answer the user query.
  *
  * @author Thomas Vitale
  * @since 1.0.0
  */
-public interface QueryAugmentor extends BiFunction<Query, List<Document>, Query> {
+public interface QueryAugmenter extends BiFunction<Query, List<Document>, Query> {
 
 	/**
 	 * Augments the user query with contextual data.
@@ -39,12 +39,6 @@ public interface QueryAugmentor extends BiFunction<Query, List<Document>, Query>
 	 */
 	Query augment(Query query, List<Document> documents);
 
-	/**
-	 * Augments the user query with contextual data.
-	 * @param query The user query to augment
-	 * @param documents The contextual data to use for augmentation
-	 * @return The augmented query
-	 */
 	default Query apply(Query query, List<Document> documents) {
 		return augment(query, documents);
 	}

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

@@ -15,11 +15,11 @@
  */
 
 /**
- * RAG Component: Query Transformation.
+ * RAG Sub-Module: Query Augmentation.
  */
 @NonNullApi
 @NonNullFields
-package org.springframework.ai.rag.analysis.query.transformation;
+package org.springframework.ai.rag.generation.augmentation;
 
 import org.springframework.lang.NonNullApi;
 import org.springframework.lang.NonNullFields;

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

@@ -15,15 +15,14 @@
  */
 
 /**
- * RAG Module: Query Analysis.
+ * RAG Module: Generation.
  * <p>
- * This package encompasses all components involved in the pre-retrieval phase of a
- * retrieval augmented generation flow. Queries are transformed, expanded, or constructed
- * so to enhance the effectiveness and accuracy of the subsequent retrieval phase.
+ * This package includes components for handling the generation stage in Retrieval
+ * Augmented Generation flows.
  */
 @NonNullApi
 @NonNullFields
-package org.springframework.ai.rag.analysis;
+package org.springframework.ai.rag.generation;
 
 import org.springframework.lang.NonNullApi;
 import org.springframework.lang.NonNullFields;

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

@@ -15,16 +15,14 @@
  */
 
 /**
- * RAG Module: Query Augmentation.
+ * RAG Module: Orchestration.
  * <p>
- * This package encompasses all components involved in the augmentation phase of a
- * retrieval augmented generation flow. The goal of this phase is to enrich the user query
- * with additional context that can be used to improve the quality of the generated
- * response.
+ * This package includes components for controlling the execution flow in a Retrieval
+ * Augmented Generation system.
  */
 @NonNullApi
 @NonNullFields
-package org.springframework.ai.rag.augmentation;
+package org.springframework.ai.rag.orchestration;
 
 import org.springframework.lang.NonNullApi;
 import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouter.java

@@ -0,0 +1,81 @@
+/*
+ * 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.orchestration.routing;
+
+import java.util.Arrays;
+import java.util.List;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
+import org.springframework.util.Assert;
+
+/**
+ * Routes a query to all the defined document retrievers.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public class AllRetrieversQueryRouter implements QueryRouter {
+
+	private static final Logger logger = LoggerFactory.getLogger(AllRetrieversQueryRouter.class);
+
+	private final List<DocumentRetriever> documentRetrievers;
+
+	public AllRetrieversQueryRouter(List<DocumentRetriever> documentRetrievers) {
+		Assert.notEmpty(documentRetrievers, "documentRetrievers cannot be null or empty");
+		Assert.noNullElements(documentRetrievers, "documentRetrievers cannot contain null elements");
+		this.documentRetrievers = documentRetrievers;
+	}
+
+	@Override
+	public List<DocumentRetriever> route(Query query) {
+		Assert.notNull(query, "query cannot be null");
+		logger.debug("Routing query to all document retrievers");
+		return this.documentRetrievers;
+	}
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public final static class Builder {
+
+		private List<DocumentRetriever> documentRetrievers;
+
+		private Builder() {
+		}
+
+		public Builder documentRetrievers(DocumentRetriever... documentRetrievers) {
+			this.documentRetrievers = Arrays.asList(documentRetrievers);
+			return this;
+		}
+
+		public Builder documentRetrievers(List<DocumentRetriever> documentRetrievers) {
+			this.documentRetrievers = documentRetrievers;
+			return this;
+		}
+
+		public AllRetrieversQueryRouter build() {
+			return new AllRetrieversQueryRouter(this.documentRetrievers);
+		}
+
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/QueryRouter.java

@@ -0,0 +1,53 @@
+/*
+ * 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.orchestration.routing;
+
+import java.util.List;
+import java.util.function.Function;
+
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.retrieval.join.DocumentJoiner;
+import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
+
+/**
+ * A component for routing a query to one or more document retrievers. It provides a
+ * decision-making mechanism to support various scenarios and making the Retrieval
+ * Augmented Generation flow more flexible and extensible. It can be used to implement
+ * routing strategies using metadata, large language models, tools (the foundation of
+ * Agentic RAG), and other techniques.
+ * <p>
+ * When retrieving documents from multiple sources, you'll need to join the results before
+ * concluding the retrieval stage. For this purpose, you can use the
+ * {@link DocumentJoiner}.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public interface QueryRouter extends Function<Query, List<DocumentRetriever>> {
+
+	/**
+	 * Routes a query to one or more document retrievers.
+	 * @param query the query to route
+	 * @return a list of document retrievers
+	 */
+	List<DocumentRetriever> route(Query query);
+
+	default List<DocumentRetriever> apply(Query query) {
+		return route(query);
+	}
+
+}

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

@@ -15,11 +15,11 @@
  */
 
 /**
- * RAG Component: Query Expansion.
+ * RAG Sub-Module: Query Router.
  */
 @NonNullApi
 @NonNullFields
-package org.springframework.ai.rag.analysis.query.expansion;
+package org.springframework.ai.rag.orchestration.routing;
 
 import org.springframework.lang.NonNullApi;
 import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/compression/DocumentCompressor.java

@@ -0,0 +1,51 @@
+/*
+ * 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.postretrieval.compression;
+
+import java.util.List;
+import java.util.function.BiFunction;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.postretrieval.ranking.DocumentRanker;
+import org.springframework.ai.rag.postretrieval.selection.DocumentSelector;
+
+/**
+ * A component for compressing the content of each document to reduce noise and redundancy
+ * in the retrieved information, addressing challenges such as "lost-in-the-middle" and
+ * context length restrictions from the model.
+ * <p>
+ * Unlike {@link DocumentSelector}, this component does not remove entire documents from
+ * the list, but rather alters the content of the documents. Unlike
+ * {@link DocumentRanker}, this component does not change the order/score of the documents
+ * in the list.
+ */
+public interface DocumentCompressor extends BiFunction<Query, List<Document>, List<Document>> {
+
+	/**
+	 * Compresses the content of each document.
+	 * @param query the query to compress documents for
+	 * @param documents the list of documents whose content should be compressed
+	 * @return a list of documents with compressed content
+	 */
+	List<Document> compress(Query query, List<Document> documents);
+
+	default List<Document> apply(Query query, List<Document> documents) {
+		return compress(query, documents);
+	}
+
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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: Document Compression.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.postretrieval.compression;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/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: Post-Retrieval.
+ * <p>
+ * This package includes components for handling the post-retrieval stage in Retrieval
+ * Augmented Generation flows.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.postretrieval;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/ranking/DocumentRanker.java

@@ -0,0 +1,50 @@
+/*
+ * 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.postretrieval.ranking;
+
+import java.util.List;
+import java.util.function.BiFunction;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.postretrieval.compression.DocumentCompressor;
+import org.springframework.ai.rag.postretrieval.selection.DocumentSelector;
+
+/**
+ * A component for ordering and ranking documents based on their relevance to a query to
+ * bring the most relevant documents to the top of the list, addressing challenges such as
+ * "lost-in-the-middle".
+ * <p>
+ * Unlike {@link DocumentSelector}, this component does not remove entire documents from
+ * the list, but rather changes the order/score of the documents in the list. Unlike
+ * {@link DocumentCompressor}, this component does not alter the content of the documents.
+ */
+public interface DocumentRanker extends BiFunction<Query, List<Document>, List<Document>> {
+
+	/**
+	 * Ranks documents based on their relevance to the given query.
+	 * @param query the query to rank documents for
+	 * @param documents the list of documents to rank
+	 * @return a list of ordered documents based on a ranking algorithm
+	 */
+	List<Document> rank(Query query, List<Document> documents);
+
+	default List<Document> apply(Query query, List<Document> documents) {
+		return rank(query, documents);
+	}
+
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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: Document Ranking.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.postretrieval.ranking;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/selection/DocumentSelector.java

@@ -0,0 +1,51 @@
+/*
+ * 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.postretrieval.selection;
+
+import java.util.List;
+import java.util.function.BiFunction;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.postretrieval.compression.DocumentCompressor;
+import org.springframework.ai.rag.postretrieval.ranking.DocumentRanker;
+
+/**
+ * A component for removing irrelevant or redundant documents from a list of retrieved
+ * documents, addressing challenges such as "lost-in-the-middle" and context length
+ * restrictions from the model.
+ * <p>
+ * Unlike {@link DocumentRanker}, this component does not change the order/score of the
+ * documents in the list, but rather removes irrelevant or redundant documents. Unlike
+ * {@link DocumentCompressor}, this component does not alter the content of the documents,
+ * but rather removes entire documents.
+ */
+public interface DocumentSelector extends BiFunction<Query, List<Document>, List<Document>> {
+
+	/**
+	 * Removes irrelevant or redundant documents from a list of retrieved documents.
+	 * @param query the query to select documents for
+	 * @param documents the list of documents to select from
+	 * @return a list of selected documents
+	 */
+	List<Document> select(Query query, List<Document> documents);
+
+	default List<Document> apply(Query query, List<Document> documents) {
+		return select(query, documents);
+	}
+
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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: Document Selection.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.postretrieval.selection;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/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: Pre-Retrieval.
+ * <p>
+ * This package includes components for handling the pre-retrieval stage in Retrieval
+ * Augmented Generation flows.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.preretrieval;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/MultiQueryExpander.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.analysis.query.expansion;
+package org.springframework.ai.rag.preretrieval.query.expansion;
 
 import java.util.Arrays;
 import java.util.List;
@@ -33,10 +33,9 @@ import org.springframework.util.CollectionUtils;
 import org.springframework.util.StringUtils;
 
 /**
- * Expander that implements semantic query expansion for retrieval-augmented generation
- * flows. It uses a large language model to generate multiple semantically diverse
- * variations of an input query to capture different perspectives and improve document
- * retrieval coverage.
+ * Uses a large language model to expand a query into multiple semantically diverse
+ * variations to capture different perspectives, useful for retrieving additional
+ * contextual information and increasing the chances of finding relevant results.
  *
  * <p>
  * Example usage: <pre>{@code
@@ -70,7 +69,7 @@ public final class MultiQueryExpander implements QueryExpander {
 			Query variants:
 			""");
 
-	private static final Boolean DEFAULT_INCLUDE_ORIGINAL = false;
+	private static final Boolean DEFAULT_INCLUDE_ORIGINAL = true;
 
 	private static final Integer DEFAULT_NUMBER_OF_QUERIES = 3;
 

spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/QueryExpander.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.analysis.query.expansion;
+package org.springframework.ai.rag.preretrieval.query.expansion;
 
 import java.util.List;
 import java.util.function.Function;
@@ -22,10 +22,9 @@ import java.util.function.Function;
 import org.springframework.ai.rag.Query;
 
 /**
- * A component responsible for expanding the input query into a list of related queries
- * based on a specified strategy. These expansions can be used to capture different
- * perspectives or to break down complex queries into simpler, more manageable
- * sub-queries, thereby improving the retrieval process.
+ * A component for expanding the input query into a list of queries, addressing challenges
+ * such as poorly formed queries by providing alternative query formulations, or by
+ * breaking down complex problems into simpler sub-queries,
  *
  * @author Thomas Vitale
  * @since 1.0.0
@@ -33,19 +32,12 @@ import org.springframework.ai.rag.Query;
 public interface QueryExpander extends Function<Query, List<Query>> {
 
 	/**
-	 * Expands the given query into a list of related queries according to the implemented
-	 * strategy.
+	 * Expands the given query into a list of queries.
 	 * @param query The original query to be expanded
 	 * @return A list of expanded queries
 	 */
 	List<Query> expand(Query query);
 
-	/**
-	 * Expands the given query into a list of related queries according to the implemented
-	 * strategy.
-	 * @param query The original query to be expanded
-	 * @return A list of expanded queries
-	 */
 	default List<Query> apply(Query query) {
 		return expand(query);
 	}

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

@@ -0,0 +1,25 @@
+/*
+ * 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: Query Expansion.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.preretrieval.query.expansion;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/QueryTransformer.java

@@ -14,16 +14,16 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.analysis.query.transformation;
+package org.springframework.ai.rag.preretrieval.query.transformation;
 
 import java.util.function.Function;
 
 import org.springframework.ai.rag.Query;
 
 /**
- * Component responsible for transforming the input query based on a specified strategy.
- * These transformations can be used to enhance the clarity, semantic meaning, or language
- * of the query, thereby improving the effectiveness of the retrieval process.
+ * A component for transforming the input query to make it more effective for retrieval
+ * tasks, addressing challenges such as poorly formed queries, ambiguous terms, complex
+ * vocabulary, or unsupported languages.
  *
  * @author Thomas Vitale
  * @since 1.0.0
@@ -37,11 +37,6 @@ public interface QueryTransformer extends Function<Query, Query> {
 	 */
 	Query transform(Query query);
 
-	/**
-	 * Transforms the given query according to the implemented strategy.
-	 * @param query The original query to transform
-	 * @return The transformed query
-	 */
 	default Query apply(Query query) {
 		return transform(query);
 	}

spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/TranslationQueryTransformer.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.analysis.query.transformation;
+package org.springframework.ai.rag.preretrieval.query.transformation;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -29,10 +29,13 @@ import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 
 /**
- * Transformer that handles translation of the input query to a target language using a
- * large language model. It's aimed at optimizing similarity searches by translating a
- * query into a language supported by the document store.
- *
+ * Uses a large language model to translate a query to a target language that is supported
+ * by the embedding model used to generate the document embeddings. If the query is
+ * already in the target language, it is returned unchanged. If the language of the query
+ * is unknown, it is also returned unchanged.
+ * <p>
+ * This transformer is useful when the embedding model is trained on a specific language
+ * and the user query is in a different language.
  * <p>
  * Example usage: <pre>{@code
  * QueryTransformer transformer = TranslationQueryTransformer.builder()

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

@@ -0,0 +1,25 @@
+/*
+ * 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: Query Transformation.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.preretrieval.query.transformation;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java

@@ -0,0 +1,60 @@
+/*
+ * 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.join;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.util.Assert;
+
+/**
+ * Combines documents retrieved based on multiple queries and from multiple data sources
+ * by concatenating them into a single collection of documents. In case of duplicate
+ * documents, the first occurrence is kept. The score of each document is kept as is.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public class ConcatenationDocumentJoiner implements DocumentJoiner {
+
+	private static final Logger logger = LoggerFactory.getLogger(ConcatenationDocumentJoiner.class);
+
+	@Override
+	public List<Document> join(Map<Query, List<List<Document>>> documentsForQuery) {
+		Assert.notNull(documentsForQuery, "documentsForQuery cannot be null");
+		Assert.noNullElements(documentsForQuery.keySet(), "documentsForQuery cannot contain null keys");
+		Assert.noNullElements(documentsForQuery.values(), "documentsForQuery cannot contain null values");
+
+		logger.debug("Joining documents by concatenation");
+
+		return new ArrayList<>(documentsForQuery.values()
+			.stream()
+			.flatMap(List::stream)
+			.flatMap(List::stream)
+			.collect(Collectors.toMap(Document::getId, Function.identity(), (existing, duplicate) -> existing))
+			.values());
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/DocumentJoiner.java

@@ -0,0 +1,48 @@
+/*
+ * 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.join;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+
+/**
+ * A component for combining documents retrieved based on multiple queries and from
+ * multiple data sources into a single collection of documents. As part of the joining
+ * process, it can also handle duplicate documents and reciprocal ranking strategies.
+ *
+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public interface DocumentJoiner extends Function<Map<Query, List<List<Document>>>, List<Document>> {
+
+	/**
+	 * Joins documents retrieved across multiple queries and daa sources.
+	 * @param documentsForQuery a map of queries and the corresponding list of documents
+	 * retrieved
+	 * @return a single collection of documents
+	 */
+	List<Document> join(Map<Query, List<List<Document>>> documentsForQuery);
+
+	default List<Document> apply(Map<Query, List<List<Document>>> documentsForQuery) {
+		return join(documentsForQuery);
+	}
+
+}

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

@@ -0,0 +1,25 @@
+/*
+ * 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: Document Join.
+ */
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.rag.retrieval.join;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;

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

@@ -17,8 +17,8 @@
 /**
  * RAG Module: Information Retrieval.
  * <p>
- * This package includes submodules for handling the retrieval process in
- * retrieval-augmented generation flows.
+ * This package includes components for handling the retrieval stage in Retrieval
+ * Augmented Generation flows.
  */
 @NonNullApi
 @NonNullFields

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

@@ -40,12 +40,6 @@ public interface DocumentRetriever extends Function<Query, List<Document>> {
 	 */
 	List<Document> retrieve(Query query);
 
-	/**
-	 * Retrieves relevant documents from an underlying data source based on the given
-	 * query.
-	 * @param query The query to use for retrieving documents
-	 * @return The list of relevant documents
-	 */
 	default List<Document> apply(Query query) {
 		return retrieve(query);
 	}

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

@@ -28,8 +28,9 @@ import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 /**
- * Document retriever that uses a vector store to search for documents. It supports
- * filtering based on metadata, similarity threshold, and top-k results.
+ * Retrieves documents from a vector store that are semantically similar to the input
+ * query. It supports filtering based on metadata, similarity threshold, and top-k
+ * results.
  *
  * <p>
  * Example usage: <pre>{@code
@@ -61,6 +62,9 @@ public final class VectorStoreDocumentRetriever implements DocumentRetriever {
 	public VectorStoreDocumentRetriever(VectorStore vectorStore, @Nullable Double similarityThreshold,
 			@Nullable Integer topK, @Nullable Supplier<Filter.Expression> filterExpression) {
 		Assert.notNull(vectorStore, "vectorStore cannot be null");
+		Assert.isTrue(similarityThreshold == null || similarityThreshold >= 0.0,
+				"similarityThreshold must be equal to or greater than 0.0");
+		Assert.isTrue(topK == null || topK > 0, "topK must be greater than 0");
 		this.vectorStore = vectorStore;
 		this.similarityThreshold = similarityThreshold != null ? similarityThreshold
 				: SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL;
@@ -104,14 +108,11 @@ public final class VectorStoreDocumentRetriever implements DocumentRetriever {
 		}
 
 		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;
 		}

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

@@ -15,7 +15,7 @@
  */
 
 /**
- * RAG Component: Document Search.
+ * RAG Sub-Module: Document Search.
  */
 @NonNullApi
 @NonNullFields

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

@@ -16,11 +16,8 @@
 
 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;
@@ -29,9 +26,11 @@ 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.analysis.query.transformation.QueryTransformer;
+import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
 import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
 
+import java.util.List;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.BDDMockito.given;
@@ -44,24 +43,6 @@ import static org.mockito.Mockito.mock;
  */
 class RetrievalAugmentationAdvisorTests {
 
-	@Test
-	void whenQueryTransformerListIsNullThenThrow() {
-		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder()
-			.queryTransformers((List<QueryTransformer>) null)
-			.documentRetriever(mock(DocumentRetriever.class))
-			.build()).isInstanceOf(IllegalArgumentException.class)
-			.hasMessageContaining("queryTransformers cannot be null");
-	}
-
-	@Test
-	void whenQueryTransformerArrayIsNullThenThrow() {
-		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder()
-			.queryTransformers((QueryTransformer[]) null)
-			.documentRetriever(mock(DocumentRetriever.class))
-			.build()).isInstanceOf(IllegalArgumentException.class)
-			.hasMessageContaining("queryTransformers cannot be null");
-	}
-
 	@Test
 	void whenQueryTransformersContainNullElementsThenThrow() {
 		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder()
@@ -72,10 +53,10 @@ class RetrievalAugmentationAdvisorTests {
 	}
 
 	@Test
-	void whenDocumentRetrieverIsNullThenThrow() {
-		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder().documentRetriever(null).build())
+	void whenQueryRouterIsNullThenThrow() {
+		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder().queryRouter(null).build())
 			.isInstanceOf(IllegalArgumentException.class)
-			.hasMessageContaining("documentRetriever cannot be null");
+			.hasMessageContaining("queryRouter cannot be null");
 	}
 
 	@Test

spring-ai-core/src/test/java/org/springframework/ai/rag/generation/augmentation/ContextualQueryAugmenterTests.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.augmentation;
+package org.springframework.ai.rag.generation.augmentation;
 
 import java.util.List;
 import java.util.Map;
@@ -29,16 +29,16 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
- * Unit tests for {@link ContextualQueryAugmentor}.
+ * Unit tests for {@link ContextualQueryAugmenter}.
  *
  * @author Thomas Vitale
  */
-class ContextualQueryAugmentorTests {
+class ContextualQueryAugmenterTests {
 
 	@Test
 	void whenPromptHasMissingContextPlaceholderThenThrow() {
 		PromptTemplate customPromptTemplate = new PromptTemplate("You are the boss. Query: {query}");
-		assertThatThrownBy(() -> ContextualQueryAugmentor.builder().promptTemplate(customPromptTemplate).build())
+		assertThatThrownBy(() -> ContextualQueryAugmenter.builder().promptTemplate(customPromptTemplate).build())
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("The following placeholders must be present in the prompt template")
 			.hasMessageContaining("context");
@@ -47,7 +47,7 @@ class ContextualQueryAugmentorTests {
 	@Test
 	void whenPromptHasMissingQueryPlaceholderThenThrow() {
 		PromptTemplate customPromptTemplate = new PromptTemplate("You are the boss. Context: {context}");
-		assertThatThrownBy(() -> ContextualQueryAugmentor.builder().promptTemplate(customPromptTemplate).build())
+		assertThatThrownBy(() -> ContextualQueryAugmenter.builder().promptTemplate(customPromptTemplate).build())
 			.isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("The following placeholders must be present in the prompt template")
 			.hasMessageContaining("query");
@@ -55,36 +55,35 @@ class ContextualQueryAugmentorTests {
 
 	@Test
 	void whenQueryIsNullThenThrow() {
-		QueryAugmentor augmenter = ContextualQueryAugmentor.builder().build();
+		QueryAugmenter augmenter = ContextualQueryAugmenter.builder().build();
 		assertThatThrownBy(() -> augmenter.augment(null, List.of())).isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("query cannot be null");
 	}
 
 	@Test
 	void whenDocumentsIsNullThenThrow() {
-		QueryAugmentor augmentor = ContextualQueryAugmentor.builder().build();
+		QueryAugmenter augmenter = ContextualQueryAugmenter.builder().build();
 		Query query = new Query("test query");
-		assertThatThrownBy(() -> augmentor.augment(query, null)).isInstanceOf(IllegalArgumentException.class)
+		assertThatThrownBy(() -> augmenter.augment(query, null)).isInstanceOf(IllegalArgumentException.class)
 			.hasMessageContaining("documents cannot be null");
 	}
 
 	@Test
 	void whenDocumentsIsEmptyAndAllowEmptyContextThenReturnOriginalQuery() {
-		QueryAugmentor augmentor = ContextualQueryAugmentor.builder().build();
+		QueryAugmenter augmenter = ContextualQueryAugmenter.builder().allowEmptyContext(true).build();
 		Query query = new Query("test query");
-		Query augmentedQuery = augmentor.augment(query, List.of());
+		Query augmentedQuery = augmenter.augment(query, List.of());
 		assertThat(augmentedQuery).isEqualTo(query);
 	}
 
 	@Test
 	void whenDocumentsIsEmptyAndNotAllowEmptyContextThenReturnAugmentedQueryWithCustomTemplate() {
 		PromptTemplate emptyContextPromptTemplate = new PromptTemplate("No context available.");
-		QueryAugmentor augmentor = ContextualQueryAugmentor.builder()
-			.allowEmptyContext(false)
+		QueryAugmenter augmenter = ContextualQueryAugmenter.builder()
 			.emptyContextPromptTemplate(emptyContextPromptTemplate)
 			.build();
 		Query query = new Query("test query");
-		Query augmentedQuery = augmentor.augment(query, List.of());
+		Query augmentedQuery = augmenter.augment(query, List.of());
 		assertThat(augmentedQuery.text()).isEqualTo(emptyContextPromptTemplate.getTemplate());
 	}
 
@@ -97,10 +96,10 @@ class ContextualQueryAugmentorTests {
 				Query:
 				{query}
 				""");
-		QueryAugmentor augmentor = ContextualQueryAugmentor.builder().promptTemplate(promptTemplate).build();
+		QueryAugmenter augmenter = ContextualQueryAugmenter.builder().promptTemplate(promptTemplate).build();
 		Query query = new Query("test query");
 		List<Document> documents = List.of(new Document("content1", Map.of()), new Document("content2", Map.of()));
-		Query augmentedQuery = augmentor.augment(query, documents);
+		Query augmentedQuery = augmenter.augment(query, documents);
 		assertThat(augmentedQuery.text()).isEqualTo("""
 				Context:
 				content1

spring-ai-core/src/test/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouterTests.java

@@ -0,0 +1,80 @@
+/*
+ * 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.orchestration.routing;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Unit tests for {@link AllRetrieversQueryRouter}.
+ *
+ * @author Thomas Vitale
+ */
+class AllRetrieversQueryRouterTests {
+
+	@Test
+	void whenDocumentRetrieversIsNullThenThrow() {
+		assertThatThrownBy(
+				() -> AllRetrieversQueryRouter.builder().documentRetrievers((List<DocumentRetriever>) null).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentRetrievers cannot be null or empty");
+	}
+
+	@Test
+	void whenDocumentRetrieversIsEmptyThenThrow() {
+		assertThatThrownBy(() -> AllRetrieversQueryRouter.builder().documentRetrievers(List.of()).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentRetrievers cannot be null or empty");
+	}
+
+	@Test
+	void whenDocumentRetrieversContainsNullKeysThenThrow() {
+		var documentRetrievers = new ArrayList<DocumentRetriever>();
+		documentRetrievers.add(null);
+		assertThatThrownBy(() -> AllRetrieversQueryRouter.builder().documentRetrievers(documentRetrievers).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentRetrievers cannot contain null elements");
+	}
+
+	@Test
+	void whenQueryIsNullThenThrow() {
+		DocumentRetriever documentRetriever = mock(DocumentRetriever.class);
+		QueryRouter queryRouter = AllRetrieversQueryRouter.builder().documentRetrievers(documentRetriever).build();
+		assertThatThrownBy(() -> queryRouter.route(null)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("query cannot be null");
+	}
+
+	@Test
+	void routeToAllRetrievers() {
+		DocumentRetriever documentRetriever1 = mock(DocumentRetriever.class);
+		DocumentRetriever documentRetriever2 = mock(DocumentRetriever.class);
+		QueryRouter queryRouter = AllRetrieversQueryRouter.builder()
+			.documentRetrievers(documentRetriever1, documentRetriever2)
+			.build();
+		List<DocumentRetriever> selectedDocumentRetrievers = queryRouter.route(new Query("test"));
+		assertThat(selectedDocumentRetrievers).containsAll(List.of(documentRetriever1, documentRetriever2));
+	}
+
+}

spring-ai-core/src/test/java/org/springframework/ai/rag/preretrieval/query/expansion/MultiQueryExpanderTests.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.analysis.query.expansion;
+package org.springframework.ai.rag.preretrieval.query.expansion;
 
 import org.junit.jupiter.api.Test;
 

spring-ai-core/src/test/java/org/springframework/ai/rag/preretrieval/query/transformation/TranslationQueryTransformerTests.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.rag.analysis.query.transformation;
+package org.springframework.ai.rag.preretrieval.query.transformation;
 
 import org.junit.jupiter.api.Test;
 

spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java

@@ -0,0 +1,95 @@
+/*
+ * 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.join;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Unit tests for {@link ConcatenationDocumentJoiner}.
+ *
+ * @author Thomas Vitale
+ */
+class ConcatenationDocumentJoinerTests {
+
+	@Test
+	void whenDocumentsForQueryIsNullThenThrow() {
+		DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
+		assertThatThrownBy(() -> documentJoiner.apply(null)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentsForQuery cannot be null");
+	}
+
+	@Test
+	void whenDocumentsForQueryContainsNullKeysThenThrow() {
+		DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
+		var documentsForQuery = new HashMap<Query, List<List<Document>>>();
+		documentsForQuery.put(null, List.of());
+		assertThatThrownBy(() -> documentJoiner.apply(documentsForQuery)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentsForQuery cannot contain null keys");
+	}
+
+	@Test
+	void whenDocumentsForQueryContainsNullValuesThenThrow() {
+		DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
+		var documentsForQuery = new HashMap<Query, List<List<Document>>>();
+		documentsForQuery.put(new Query("test"), null);
+		assertThatThrownBy(() -> documentJoiner.apply(documentsForQuery)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("documentsForQuery cannot contain null values");
+	}
+
+	@Test
+	void whenNoDuplicatedDocumentsThenAllDocumentsAreJoined() {
+		DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
+		var documentsForQuery = new HashMap<Query, List<List<Document>>>();
+		documentsForQuery.put(new Query("query1"),
+				List.of(List.of(new Document("1", "Content 1", Map.of()), new Document("2", "Content 2", Map.of())),
+						List.of(new Document("3", "Content 3", Map.of()))));
+		documentsForQuery.put(new Query("query2"), List.of(List.of(new Document("4", "Content 4", Map.of()))));
+
+		List<Document> result = documentJoiner.join(documentsForQuery);
+
+		assertThat(result).hasSize(4);
+		assertThat(result).extracting(Document::getId).containsExactlyInAnyOrder("1", "2", "3", "4");
+	}
+
+	@Test
+	void whenDuplicatedDocumentsThenOnlyFirstOccurrenceIsKept() {
+		DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
+		var documentsForQuery = new HashMap<Query, List<List<Document>>>();
+		documentsForQuery.put(new Query("query1"),
+				List.of(List.of(new Document("1", "Content 1", Map.of()), new Document("2", "Content 2", Map.of())),
+						List.of(new Document("3", "Content 3", Map.of()))));
+		documentsForQuery.put(new Query("query2"), List
+			.of(List.of(new Document("2", "Content 2 Duplicate", Map.of()), new Document("4", "Content 4", Map.of()))));
+
+		List<Document> result = documentJoiner.join(documentsForQuery);
+
+		assertThat(result).hasSize(4);
+		assertThat(result).extracting(Document::getId).containsExactlyInAnyOrder("1", "2", "3", "4");
+		assertThat(result).extracting(Document::getContent).contains("Content 2");
+		assertThat(result).extracting(Document::getContent).doesNotContain("Content 2 Duplicate");
+	}
+
+}

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

@@ -53,6 +53,31 @@ class VectorStoreDocumentRetrieverTests {
 			.hasMessageContaining("vectorStore cannot be null");
 	}
 
+	@Test
+	void whenTopKIsZeroThenThrow() {
+		assertThatThrownBy(
+				() -> VectorStoreDocumentRetriever.builder().topK(0).vectorStore(mock(VectorStore.class)).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("topK must be greater than 0");
+	}
+
+	@Test
+	void whenTopKIsNegativeThenThrow() {
+		assertThatThrownBy(
+				() -> VectorStoreDocumentRetriever.builder().topK(-1).vectorStore(mock(VectorStore.class)).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("topK must be greater than 0");
+	}
+
+	@Test
+	void whenSimilarityThresholdIsNegativeThenThrow() {
+		assertThatThrownBy(() -> VectorStoreDocumentRetriever.builder()
+			.similarityThreshold(-1.0)
+			.vectorStore(mock(VectorStore.class))
+			.build()).isInstanceOf(IllegalArgumentException.class)
+			.hasMessageContaining("similarityThreshold must be equal to or greater than 0.0");
+	}
+
 	@Test
 	void searchRequestParameters() {
 		var mockVectorStore = mock(VectorStore.class);

spring-ai-integration-tests/pom.xml

@@ -48,6 +48,13 @@
 			<scope>test</scope>
 		</dependency>
 
+		<dependency>
+			<groupId>io.micrometer</groupId>
+			<artifactId>context-propagation</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+
 		<dependency>
 			<groupId>org.springframework.ai</groupId>
 			<artifactId>spring-ai-openai-spring-boot-starter</artifactId>

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

@@ -33,7 +33,8 @@ import org.springframework.ai.evaluation.EvaluationResponse;
 import org.springframework.ai.evaluation.RelevancyEvaluator;
 import org.springframework.ai.integration.tests.TestApplication;
 import org.springframework.ai.openai.OpenAiChatModel;
-import org.springframework.ai.rag.analysis.query.transformation.TranslationQueryTransformer;
+import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
+import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
 import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
 import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
 import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
@@ -114,6 +115,36 @@ class RetrievalAugmentationAdvisorIT {
 			.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
 			.build();
 
+		ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel)
+			.build()
+			.prompt()
+			.system("Answer the question in English")
+			.user(question)
+			.advisors(ragAdvisor)
+			.call()
+			.chatResponse();
+
+		assertThat(chatResponse).isNotNull();
+
+		String response = chatResponse.getResult().getOutput().getContent();
+		System.out.println(response);
+		assertThat(response).containsIgnoringCase("Highlands");
+
+		evaluateRelevancy(question, chatResponse);
+	}
+
+	@Test
+	void ragWithMultiQuery() {
+		String question = "Where does the adventure of Anacletus and Birba take place?";
+
+		RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
+			.queryExpander(MultiQueryExpander.builder()
+				.chatClientBuilder(ChatClient.builder(this.openAiChatModel))
+				.numberOfQueries(2)
+				.build())
+			.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
+			.build();
+
 		ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel)
 			.build()
 			.prompt(question)

spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/generation/augmentation/ContextualQueryAugmenterIT.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.integration.tests.rag.augmentation;
+package org.springframework.ai.integration.tests.rag.generation.augmentation;
 
 import java.util.List;
 
@@ -25,34 +25,34 @@ import org.springframework.ai.document.Document;
 import org.springframework.ai.integration.tests.TestApplication;
 import org.springframework.ai.openai.OpenAiChatModel;
 import org.springframework.ai.rag.Query;
-import org.springframework.ai.rag.augmentation.ContextualQueryAugmentor;
-import org.springframework.ai.rag.augmentation.QueryAugmentor;
+import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
+import org.springframework.ai.rag.generation.augmentation.QueryAugmenter;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
- * Integration tests for {@link ContextualQueryAugmentor}.
+ * Integration tests for {@link ContextualQueryAugmenter}.
  *
  * @author Thomas Vitale
  */
 @SpringBootTest(classes = TestApplication.class)
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
-class ContextualQueryAugmentorIT {
+class ContextualQueryAugmenterIT {
 
 	@Autowired
 	OpenAiChatModel openAiChatModel;
 
 	@Test
 	void whenContextIsProvided() {
-		QueryAugmentor queryAugmentor = ContextualQueryAugmentor.builder().build();
+		QueryAugmenter queryAugmenter = ContextualQueryAugmenter.builder().build();
 		Query query = new Query("What is Iorek's dream?");
 		List<Document> documents = List
 			.of(new Document("Iorek was a little polar bear who lived in the Arctic circle."), new Document(
 					"Iorek loved to explore the snowy landscape and dreamt of one day going on an adventure around the North Pole."));
 
-		Query augmentedQuery = queryAugmentor.augment(query, documents);
+		Query augmentedQuery = queryAugmenter.augment(query, documents);
 		String response = this.openAiChatModel.call(augmentedQuery.text());
 
 		assertThat(response).isNotEmpty();
@@ -64,10 +64,10 @@ class ContextualQueryAugmentorIT {
 
 	@Test
 	void whenAllowEmptyContext() {
-		QueryAugmentor queryAugmentor = ContextualQueryAugmentor.builder().build();
+		QueryAugmenter queryAugmenter = ContextualQueryAugmenter.builder().allowEmptyContext(true).build();
 		Query query = new Query("What is Iorek's dream?");
 		List<Document> documents = List.of();
-		Query augmentedQuery = queryAugmentor.augment(query, documents);
+		Query augmentedQuery = queryAugmenter.augment(query, documents);
 		String response = this.openAiChatModel.call(augmentedQuery.text());
 
 		assertThat(response).isNotEmpty();
@@ -77,10 +77,10 @@ class ContextualQueryAugmentorIT {
 
 	@Test
 	void whenNotAllowEmptyContext() {
-		QueryAugmentor queryAugmentor = ContextualQueryAugmentor.builder().allowEmptyContext(false).build();
+		QueryAugmenter queryAugmenter = ContextualQueryAugmenter.builder().build();
 		Query query = new Query("What is Iorek's dream?");
 		List<Document> documents = List.of();
-		Query augmentedQuery = queryAugmentor.augment(query, documents);
+		Query augmentedQuery = queryAugmenter.augment(query, documents);
 		String response = this.openAiChatModel.call(augmentedQuery.text());
 
 		assertThat(response).isNotEmpty();

spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/expansion/MultiQueryExpanderIT.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.integration.tests.rag.analysis.query.expansion;
+package org.springframework.ai.integration.tests.rag.preretrieval.query.expansion;
 
 import java.util.List;
 
@@ -25,8 +25,8 @@ import org.springframework.ai.chat.client.ChatClient;
 import org.springframework.ai.integration.tests.TestApplication;
 import org.springframework.ai.openai.OpenAiChatModel;
 import org.springframework.ai.rag.Query;
-import org.springframework.ai.rag.analysis.query.expansion.MultiQueryExpander;
-import org.springframework.ai.rag.analysis.query.expansion.QueryExpander;
+import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
+import org.springframework.ai.rag.preretrieval.query.expansion.QueryExpander;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 
@@ -55,7 +55,7 @@ class MultiQueryExpanderIT {
 
 		assertThat(queries).isNotNull();
 		queries.forEach(System.out::println);
-		assertThat(queries).hasSize(3);
+		assertThat(queries).hasSize(4);
 	}
 
 	@Test
@@ -70,23 +70,23 @@ class MultiQueryExpanderIT {
 
 		assertThat(queries).isNotNull();
 		queries.forEach(System.out::println);
-		assertThat(queries).hasSize(4);
+		assertThat(queries).hasSize(5);
 	}
 
 	@Test
-	void whenExpanderWithOriginalQueryIncluded() {
+	void whenExpanderWithoutOriginalQueryIncluded() {
 		Query query = new Query("What is the weather in Rome?");
 		QueryExpander queryExpander = MultiQueryExpander.builder()
 			.chatClientBuilder(ChatClient.builder(this.openAiChatModel))
 			.numberOfQueries(3)
-			.includeOriginal(true)
+			.includeOriginal(false)
 			.build();
 
 		List<Query> queries = queryExpander.apply(query);
 
 		assertThat(queries).isNotNull();
 		queries.forEach(System.out::println);
-		assertThat(queries).hasSize(4);
+		assertThat(queries).hasSize(3);
 	}
 
 }

spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/TranslationQueryTransformerIT.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.ai.integration.tests.rag.analysis.query.transformation;
+package org.springframework.ai.integration.tests.rag.preretrieval.query.transformation;
 
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -23,8 +23,8 @@ import org.springframework.ai.chat.client.ChatClient;
 import org.springframework.ai.integration.tests.TestApplication;
 import org.springframework.ai.openai.OpenAiChatModel;
 import org.springframework.ai.rag.Query;
-import org.springframework.ai.rag.analysis.query.transformation.QueryTransformer;
-import org.springframework.ai.rag.analysis.query.transformation.TranslationQueryTransformer;
+import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
+import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;