Skip to main content

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) {}

Modular RAG Flow​

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();
Last updated on