Modular RAG

The spring-ai-rag module in Spring AI provides support for modular RAG. This module defines common interfaces for different phases of modular RAG. It also provides RetrievalAugmentationAdvisor as the entry point of using modular RAG.

Query

Query represents a query in a RAG flow. Query is a record type with a builder to create objects.

Query

public record Query(
    String text,
    List<Message> history,
    Map<String, Object> context) {}

Pre-Retrieval

QueryTransformer

QueryTransformer transforms the input query to make it more effective for retrieval tasks.

QueryTransformer

public interface QueryTransformer extends Function<Query, Query> {}

Spring AI provides the following implementations for QueryTransformer.

• RewriteQueryTransformer uses an LLM to rewrite a user query to provide better results when querying a target system.
• CompressionQueryTransformer uses an LLM to compress a conversation history and a follow-up query into a standalone query that captures the essence of the conversation.
• TranslationQueryTransformer uses an LLM to translate a query to a target language that is supported by the embedding model used to generate the document embeddings.

QueryExpander

QueryExpander expands the input query into a list of queries.

QueryExpander

public interface QueryExpander extends Function<Query, List<Query>> {}

Spring AI provides the following implementation of QueryExpander.

• MultiQueryExpander uses an LLM to expand a query into multiple semantically diverse variations to capture different perspectives.

Retrieval

DocumentRetriever

DocumentRetriever retrieves Documents from a data source.

DocumentRetriever

public interface DocumentRetriever extends Function<Query, List<Document>> {}

VectorStoreDocumentRetriever

VectorStoreDocumentRetriever is a built-in implementation of DocumentRetriever that retrieves documents from a vector store. A VectorStore is required when creating a VectorStoreDocumentRetriever.

Create VectorStoreDocumentRetriever

VectorStoreDocumentRetriever.builder()
    .vectorStore(vectorStore)
    .build();

DocumentJoiner

DocumentJoiner combines documents retrieved based on multiple queries and from multiple data sources into a single collection of documents.

DocumentJoiner

public interface DocumentJoiner extends Function<Map<Query, List<List<Document>>>, List<Document>>

Spring AI provides the following implementation of DocumentJoiner.

• ConcatenationDocumentJoiner combines documents retrieved based on multiple queries and from multiple data sources by concatenating them into a single collection of documents.

Post-Retrieval

DocumentPostProcessor post-processes retrieved documents based on a query.

DocumentPostProcessor

public interface DocumentPostProcessor extends BiFunction<Query, List<Document>, List<Document>> {}

Generation

QueryAugmenter augments an input query with additional data.

QueryAugmenter

public interface QueryAugmenter extends BiFunction<Query, List<Document>, Query> {}

Spring AI provides the following implementation of QueryAugmenter.

• ContextualQueryAugmenter augments the user query with contextual data from the content of the provided documents.

RetrievalAugmentationAdvisor

RetrievalAugmentationAdvisor is an Advisor that implements modular RAG flow. When creating a RetrievalAugmentationAdvisor, interfaces mentioned above can be used. Only a DocumentRetriever is required. RetrievalAugmentationAdvisors are created using a builder.

The code below creates a RetrievalAugmentationAdvisor using only a VectorStoreDocumentRetriever. This is an implementation of Naive RAG.

Create RetrievalAugmentationAdvisor

RetrievalAugmentationAdvisor.builder()
    .documentRetriever(
        VectorStoreDocumentRetriever.builder()
            .vectorStore(vectorStore)
            .build())
    .build();