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