diff --git a/spring-ai-core/pom.xml b/spring-ai-core/pom.xml index e5fd58c5782..c1bcb0f3bd4 100644 --- a/spring-ai-core/pom.xml +++ b/spring-ai-core/pom.xml @@ -96,6 +96,11 @@ micrometer-core + + io.micrometer + context-propagation + + io.micrometer micrometer-tracing-bridge-otel @@ -195,4 +200,4 @@ - \ No newline at end of file + diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java index c787bb5bca8..557de9ece55 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java @@ -16,41 +16,40 @@ 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 reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; +import java.util.concurrent.CompletableFuture; +import java.util.stream.Collectors; 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; +import reactor.core.scheduler.Scheduler; /** * 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. - *

- * It's the successor of the {@link QuestionAnswerAdvisor}. * * @author Christian Tzolov * @author Thomas Vitale @@ -58,29 +57,40 @@ * @see arXiv:2407.21059 * @see arXiv:2312.10997 */ -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 queryTransformers; - private final DocumentRetriever documentRetriever; + @Nullable + private final QueryExpander queryExpander; + + private final QueryRouter queryRouter; + + private final DocumentJoiner documentJoiner; + + private final QueryAugmenter queryAugmenter; - private final QueryAugmentor queryAugmentor; + private final TaskExecutor taskExecutor; - private final boolean protectFromBlocking; + private final Scheduler scheduler; private final int order; - public RetrievalAugmentationAdvisor(List queryTransformers, DocumentRetriever documentRetriever, - @Nullable QueryAugmentor queryAugmentor, @Nullable Boolean protectFromBlocking, @Nullable Integer order) { - Assert.notNull(queryTransformers, "queryTransformers cannot be null"); + public RetrievalAugmentationAdvisor(@Nullable List 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 +99,7 @@ public static Builder 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 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 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 context = new HashMap<>(request.adviseContext()); // 0. Create a query from the user text and parameters. @@ -135,17 +111,47 @@ private AdvisedRequest before(AdvisedRequest request) { transformedQuery = queryTransformer.apply(transformedQuery); } - // 2. Retrieve similar documents for the original query. - List documents = this.documentRetriever.retrieve(transformedQuery); + // 2. Expand query into one or multiple queries. + List expandedQueries = queryExpander != null ? queryExpander.expand(transformedQuery) + : List.of(transformedQuery); + + // 3. Get similar documents for each query. + Map>> documentsForQuery = expandedQueries.stream() + .map(query -> CompletableFuture.supplyAsync(() -> getDocumentsForQuery(query), 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 documents = 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 = 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>> getDocumentsForQuery(Query query) { + List retrievers = queryRouter.route(query); + List> documents = retrievers.stream() + .map(retriever -> CompletableFuture.supplyAsync(() -> retriever.retrieve(query), 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 +163,9 @@ private AdvisedResponse after(AdvisedResponse advisedResponse) { return new AdvisedResponse(chatResponseBuilder.build(), advisedResponse.adviseContext()); } - private Predicate 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 scheduler; } @Override @@ -178,15 +173,31 @@ public int getOrder() { 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 queryTransformers = new ArrayList<>(); + private List queryTransformers; + + private QueryExpander queryExpander; + + private QueryRouter queryRouter; + + private DocumentJoiner documentJoiner; - private DocumentRetriever documentRetriever; + private QueryAugmenter queryAugmenter; - private QueryAugmentor queryAugmentor; + private TaskExecutor taskExecutor; - private Boolean protectFromBlocking; + private Scheduler scheduler; private Integer order; @@ -194,29 +205,49 @@ private Builder() { } public Builder queryTransformers(List 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 documentJoiner(DocumentJoiner documentJoiner) { + this.documentJoiner = documentJoiner; + return this; + } + + public Builder queryAugmenter(QueryAugmenter queryAugmenter) { + this.queryAugmenter = queryAugmenter; return this; } - public Builder queryAugmentor(QueryAugmentor queryAugmentor) { - this.queryAugmentor = queryAugmentor; + public Builder taskExecutor(TaskExecutor taskExecutor) { + this.taskExecutor = taskExecutor; return this; } - public Builder protectFromBlocking(Boolean protectFromBlocking) { - this.protectFromBlocking = protectFromBlocking; + public Builder scheduler(Scheduler scheduler) { + this.scheduler = scheduler; return this; } @@ -226,8 +257,8 @@ public Builder order(Integer order) { } 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); } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/BaseAdvisor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/BaseAdvisor.java new file mode 100644 index 00000000000..4878b45e7f2 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/BaseAdvisor.java @@ -0,0 +1,107 @@ +/* + * 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 org.springframework.ai.chat.model.ChatResponse; +import org.springframework.util.Assert; +import org.springframework.util.StringUtils; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Scheduler; +import reactor.core.scheduler.Schedulers; + +import java.util.function.Predicate; + +/** + * 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 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 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 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; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/ContextualQueryAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/ContextualQueryAugmenter.java similarity index 88% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/ContextualQueryAugmentor.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/ContextualQueryAugmenter.java index 1f929ed97dd..e2a89aa7c66 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/ContextualQueryAugmentor.java +++ b/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 @@ * *

* Example usage:

{@code
- * QueryAugmentor augmentor = ContextualQueryAugmentor.builder()
+ * QueryAugmenter augmenter = ContextualQueryAugmenter.builder()
  *    .allowEmptyContext(false)
  *    .build();
- * Query augmentedQuery = augmentor.augment(query, documents);
+ * Query augmentedQuery = augmenter.augment(query, documents);
  * }
* * @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 Query augment(Query query, List documents) { 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 Builder allowEmptyContext(Boolean allowEmptyContext) { return this; } - public ContextualQueryAugmentor build() { - return new ContextualQueryAugmentor(this.promptTemplate, this.emptyContextPromptTemplate, + public ContextualQueryAugmenter build() { + return new ContextualQueryAugmenter(this.promptTemplate, this.emptyContextPromptTemplate, this.allowEmptyContext); } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/QueryAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/QueryAugmenter.java similarity index 72% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/QueryAugmentor.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/QueryAugmenter.java index d7359f69dde..eac75b58f6b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/QueryAugmentor.java +++ b/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.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> { +public interface QueryAugmenter extends BiFunction, Query> { /** * Augments the user query with contextual data. @@ -39,12 +39,6 @@ public interface QueryAugmentor extends BiFunction, Query> */ Query augment(Query query, List 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 documents) { return augment(query, documents); } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/package-info.java similarity index 87% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/package-info.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/generation/augmentation/package-info.java index a71c508bc74..10a06948355 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/package-info.java +++ b/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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/generation/package-info.java similarity index 69% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/package-info.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/generation/package-info.java index ee9deac32ac..b59411467d7 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/package-info.java +++ b/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. *

- * 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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/package-info.java similarity index 68% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/package-info.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/package-info.java index ededce78e83..7ef0db0979e 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/augmentation/package-info.java +++ b/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. *

- * 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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouter.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouter.java new file mode 100644 index 00000000000..a87b577c752 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouter.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.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; + +import java.util.Arrays; +import java.util.List; + +/** + * 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 documentRetrievers; + + public AllRetrieversQueryRouter(List documentRetrievers) { + Assert.notEmpty(documentRetrievers, "documentRetrievers cannot be null or empty"); + Assert.noNullElements(documentRetrievers, "documentRetrievers cannot contain null elements"); + this.documentRetrievers = documentRetrievers; + } + + @Override + public List route(Query query) { + Assert.notNull(query, "query cannot be null"); + logger.debug("Routing query to all document retrievers"); + return documentRetrievers; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + + private List documentRetrievers; + + private Builder() { + } + + public Builder documentRetrievers(DocumentRetriever... documentRetrievers) { + this.documentRetrievers = Arrays.asList(documentRetrievers); + return this; + } + + public Builder documentRetrievers(List documentRetrievers) { + this.documentRetrievers = documentRetrievers; + return this; + } + + public AllRetrieversQueryRouter build() { + return new AllRetrieversQueryRouter(documentRetrievers); + } + + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/QueryRouter.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/QueryRouter.java new file mode 100644 index 00000000000..1ac34f8f118 --- /dev/null +++ b/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 org.springframework.ai.rag.Query; +import org.springframework.ai.rag.retrieval.join.DocumentJoiner; +import org.springframework.ai.rag.retrieval.search.DocumentRetriever; + +import java.util.List; +import java.util.function.Function; + +/** + * 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. + *

+ * 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> { + + /** + * Routes a query to one or more document retrievers. + * @param query the query to route + * @return a list of document retrievers + */ + List route(Query query); + + default List apply(Query query) { + return route(query); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/package-info.java similarity index 88% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/package-info.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/orchestration/routing/package-info.java index b18934d7216..59a8a597f40 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/package-info.java +++ b/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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/compression/DocumentCompressor.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/compression/DocumentCompressor.java new file mode 100644 index 00000000000..2a541d3bba1 --- /dev/null +++ b/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 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; + +import java.util.List; +import java.util.function.BiFunction; + +/** + * 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. + *

+ * 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, List> { + + /** + * 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 compress(Query query, List documents); + + default List apply(Query query, List documents) { + return compress(query, documents); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/compression/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/compression/package-info.java new file mode 100644 index 00000000000..c29363101ef --- /dev/null +++ b/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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/package-info.java new file mode 100644 index 00000000000..fe1bc010906 --- /dev/null +++ b/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. + *

+ * 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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/ranking/DocumentRanker.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/ranking/DocumentRanker.java new file mode 100644 index 00000000000..e3c089839a3 --- /dev/null +++ b/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 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; + +import java.util.List; +import java.util.function.BiFunction; + +/** + * 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". + *

+ * 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, List> { + + /** + * 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 rank(Query query, List documents); + + default List apply(Query query, List documents) { + return rank(query, documents); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/ranking/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/ranking/package-info.java new file mode 100644 index 00000000000..21838ac0331 --- /dev/null +++ b/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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/selection/DocumentSelector.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/selection/DocumentSelector.java new file mode 100644 index 00000000000..7d490f9ee32 --- /dev/null +++ b/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 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; + +import java.util.List; +import java.util.function.BiFunction; + +/** + * 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. + *

+ * 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, List> { + + /** + * 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 select(Query query, List documents); + + default List apply(Query query, List documents) { + return select(query, documents); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/selection/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/postretrieval/selection/package-info.java new file mode 100644 index 00000000000..8ebc568088a --- /dev/null +++ b/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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/package-info.java new file mode 100644 index 00000000000..cbc1dab7869 --- /dev/null +++ b/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. + *

+ * 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; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/MultiQueryExpander.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/MultiQueryExpander.java similarity index 92% rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/MultiQueryExpander.java rename to spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/MultiQueryExpander.java index e5bfc56f096..08f9f874b52 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/MultiQueryExpander.java +++ b/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.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. * *

* Example usage:

{@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;
 
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/QueryExpander.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/QueryExpander.java
similarity index 61%
rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/QueryExpander.java
rename to spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/QueryExpander.java
index 3b1b1b8f60a..bb1d5d44ef5 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/QueryExpander.java
+++ b/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 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 @@
 public interface QueryExpander extends Function> {
 
 	/**
-	 * 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 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 apply(Query query) {
 		return expand(query);
 	}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/expansion/package-info.java
new file mode 100644
index 00000000000..5f9a0a1b87f
--- /dev/null
+++ b/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;
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/QueryTransformer.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/QueryTransformer.java
similarity index 69%
rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/QueryTransformer.java
rename to spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/QueryTransformer.java
index 07d1d3b4a37..511e4750c29 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/QueryTransformer.java
+++ b/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 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);
 	}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/TranslationQueryTransformer.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/TranslationQueryTransformer.java
similarity index 88%
rename from spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/TranslationQueryTransformer.java
rename to spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/TranslationQueryTransformer.java
index 6eaeb22f7e2..4fe0837ea30 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/rag/analysis/query/transformation/TranslationQueryTransformer.java
+++ b/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.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.
+ * 

+ * This transformer is useful when the embedding model is trained on a specific language + * and the user query is in a different language. *

* Example usage:

{@code
  * QueryTransformer transformer = TranslationQueryTransformer.builder()
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/preretrieval/query/transformation/package-info.java
new file mode 100644
index 00000000000..9f714c4028c
--- /dev/null
+++ b/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;
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java
new file mode 100644
index 00000000000..8f18e760dd2
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java
@@ -0,0 +1,59 @@
+/*
+ * 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.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.util.Assert;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * 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 join(Map>> 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());
+	}
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/DocumentJoiner.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/DocumentJoiner.java
new file mode 100644
index 00000000000..80593d9dcd2
--- /dev/null
+++ b/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 org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+
+/**
+ * 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>>, List> {
+
+	/**
+	 * 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 join(Map>> documentsForQuery);
+
+	default List apply(Map>> documentsForQuery) {
+		return join(documentsForQuery);
+	}
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/package-info.java
new file mode 100644
index 00000000000..c3f1483b337
--- /dev/null
+++ b/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;
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/package-info.java
index 87af7f55a7c..cc6e0613b4c 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/package-info.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/package-info.java
@@ -17,8 +17,8 @@
 /**
  * RAG Module: Information Retrieval.
  * 

- * 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 diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/DocumentRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/DocumentRetriever.java index 7982dfe4c97..1074c50367a 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/DocumentRetriever.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/DocumentRetriever.java @@ -40,12 +40,6 @@ public interface DocumentRetriever extends Function> { */ List 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 apply(Query query) { return retrieve(query); } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java index 3fd6aa074a4..dcd226c306e 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetriever.java @@ -28,8 +28,9 @@ 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. * *

* Example usage:

{@code
@@ -61,6 +62,9 @@ public final class VectorStoreDocumentRetriever implements DocumentRetriever {
 	public VectorStoreDocumentRetriever(VectorStore vectorStore, @Nullable Double similarityThreshold,
 			@Nullable Integer topK, @Nullable Supplier 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 Builder vectorStore(VectorStore vectorStore) {
 		}
 
 		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;
 		}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/package-info.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/package-info.java
index 961f18ddae4..a1e6ba54c9b 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/search/package-info.java
+++ b/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
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisorTests.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisorTests.java
index 220bb0a9477..fc75bb437ed 100644
--- a/spring-ai-core/src/test/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisorTests.java
+++ b/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.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 @@
  */
 class RetrievalAugmentationAdvisorTests {
 
-	@Test
-	void whenQueryTransformerListIsNullThenThrow() {
-		assertThatThrownBy(() -> RetrievalAugmentationAdvisor.builder()
-			.queryTransformers((List) 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 @@ void whenQueryTransformersContainNullElementsThenThrow() {
 	}
 
 	@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
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/augmentation/ContextualQueryAugmentorTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/generation/augmentation/ContextualQueryAugmenterTests.java
similarity index 78%
rename from spring-ai-core/src/test/java/org/springframework/ai/rag/augmentation/ContextualQueryAugmentorTests.java
rename to spring-ai-core/src/test/java/org/springframework/ai/rag/generation/augmentation/ContextualQueryAugmenterTests.java
index 60d3ac278c2..1b3897ab3ed 100644
--- a/spring-ai-core/src/test/java/org/springframework/ai/rag/augmentation/ContextualQueryAugmentorTests.java
+++ b/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.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 @@ void whenPromptHasMissingContextPlaceholderThenThrow() {
 	@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 @@ void whenPromptHasMissingQueryPlaceholderThenThrow() {
 
 	@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 @@ void whenDocumentsAreProvidedThenReturnAugmentedQueryWithCustomTemplate() {
 				Query:
 				{query}
 				""");
-		QueryAugmentor augmentor = ContextualQueryAugmentor.builder().promptTemplate(promptTemplate).build();
+		QueryAugmenter augmenter = ContextualQueryAugmenter.builder().promptTemplate(promptTemplate).build();
 		Query query = new Query("test query");
 		List 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
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouterTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/orchestration/routing/AllRetrieversQueryRouterTests.java
new file mode 100644
index 00000000000..32fc82fe2ec
--- /dev/null
+++ b/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) 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();
+		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 selectedDocumentRetrievers = queryRouter.route(new Query("test"));
+		assertThat(selectedDocumentRetrievers).containsAll(List.of(documentRetriever1, documentRetriever2));
+	}
+
+}
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/analysis/query/expansion/MultiQueryExpanderTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/preretrieval/query/expansion/MultiQueryExpanderTests.java
similarity index 97%
rename from spring-ai-core/src/test/java/org/springframework/ai/rag/analysis/query/expansion/MultiQueryExpanderTests.java
rename to spring-ai-core/src/test/java/org/springframework/ai/rag/preretrieval/query/expansion/MultiQueryExpanderTests.java
index e4c094e9550..b8599de3c16 100644
--- a/spring-ai-core/src/test/java/org/springframework/ai/rag/analysis/query/expansion/MultiQueryExpanderTests.java
+++ b/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;
 
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/analysis/query/transformation/TranslationQueryTransformerTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/preretrieval/query/transformation/TranslationQueryTransformerTests.java
similarity index 97%
rename from spring-ai-core/src/test/java/org/springframework/ai/rag/analysis/query/transformation/TranslationQueryTransformerTests.java
rename to spring-ai-core/src/test/java/org/springframework/ai/rag/preretrieval/query/transformation/TranslationQueryTransformerTests.java
index e9d268b3486..f39aeea7e15 100644
--- a/spring-ai-core/src/test/java/org/springframework/ai/rag/analysis/query/transformation/TranslationQueryTransformerTests.java
+++ b/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;
 
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java
new file mode 100644
index 00000000000..a57566d81ed
--- /dev/null
+++ b/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>>();
+		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>>();
+		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>>();
+		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 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>>();
+		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 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");
+	}
+
+}
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java
index 645bd51cddf..9cb4ea0438d 100644
--- a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java
+++ b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/search/VectorStoreDocumentRetrieverTests.java
@@ -53,6 +53,31 @@ void whenVectorStoreIsNullThenThrow() {
 			.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);
diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java
index 9d4ce55f1d4..65730dbb4d9 100644
--- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java
+++ b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java
@@ -33,7 +33,8 @@
 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 @@ void ragWithTranslation() {
 			.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)
diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/augmentation/ContextualQueryAugmentorIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/generation/augmentation/ContextualQueryAugmenterIT.java
similarity index 78%
rename from spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/augmentation/ContextualQueryAugmentorIT.java
rename to spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/generation/augmentation/ContextualQueryAugmenterIT.java
index 2d092abca3e..424a17b3ff2 100644
--- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/augmentation/ContextualQueryAugmentorIT.java
+++ b/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.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 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 @@ void whenContextIsProvided() {
 
 	@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 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 @@ void whenAllowEmptyContext() {
 
 	@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 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();
diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/analysis/query/expansion/MultiQueryExpanderIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/expansion/MultiQueryExpanderIT.java
similarity index 88%
rename from spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/analysis/query/expansion/MultiQueryExpanderIT.java
rename to spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/expansion/MultiQueryExpanderIT.java
index acb016d0c2b..889039dd61b 100644
--- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/analysis/query/expansion/MultiQueryExpanderIT.java
+++ b/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.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 @@ void whenExpanderWithDefaults() {
 
 		assertThat(queries).isNotNull();
 		queries.forEach(System.out::println);
-		assertThat(queries).hasSize(3);
+		assertThat(queries).hasSize(4);
 	}
 
 	@Test
@@ -70,23 +70,23 @@ void whenExpanderWithCustomQueryNumber() {
 
 		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 queries = queryExpander.apply(query);
 
 		assertThat(queries).isNotNull();
 		queries.forEach(System.out::println);
-		assertThat(queries).hasSize(4);
+		assertThat(queries).hasSize(3);
 	}
 
 }
diff --git a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/analysis/query/transformation/TranslationQueryTransformerIT.java b/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/TranslationQueryTransformerIT.java
similarity index 87%
rename from spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/analysis/query/transformation/TranslationQueryTransformerIT.java
rename to spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/TranslationQueryTransformerIT.java
index b1953d2780d..e13d5768f5a 100644
--- a/spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/analysis/query/transformation/TranslationQueryTransformerIT.java
+++ b/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.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;