Query Transformation in Spring AI Modular RAG

1. Overview

Retrieval-Augmented Generation (RAG) queries relevant documents from a knowledge base, such as a vector database, and then feeds them to the LLM as part of the prompt so it can generate a more accurate response.

Spring AI provides a modular RAG architecture that divides RAG into four stages:

• Pre-retrieval,
• Retrieval,
• Post-retrieval, and
• Generation

In this tutorial, we’ll focus on query transformation in the pre-retrieval stage that processes the user query before retrieval.

2. Why Transform Queries?

Retrieving from a vector store is basically a similarity search over the datastore embeddings, with the query embedding as the search key.

Therefore, the relevance of the retrieved information depends on the query wording. For example, a query may be an inefficient search key due to:

• Lack of clarity – a query may be ambiguous or too verbose and mention many concepts, many irrelevant to the search. In both cases, the query’s embedding is less aligned with relevant documents.
• Language mismatch – a query in a different language from the stored documents may not produce a match even if the content is there. This gets worse if the embedding model wasn’t trained on the documents in the query’s language.
• Missing context – A follow-up question within a long conversation, such as “And how does it support RAG?”, relies on previous messages. The user knows what “it” refers to, but the RAG system searches with the latest message only. Without the chat history, the query is ambiguous.

Query transformers use an LLM to rewrite queries before retrieval in a way we expect to improve the relevance of the retrieved data.

3. Setup

Let’s set up our project. First, we need the Spring AI RAG dependency to work with modular RAG:

<dependency>
   <groupId>org.springframework.ai</groupId>
   <artifactId>spring-ai-rag</artifactId>
</dependency>

As every query transformer needs a chat model that connects to an LLM, we’ll use the Spring AI OpenAI dependency:

<dependency>
   <groupId>org.springframework.ai</groupId>
   <artifactId>spring-ai-starter-model-openai</artifactId>
</dependency>

Now, let’s configure our OpenAI API key and the chat model in our application.yml:

spring:
  ai:
    openai:
      api-key: ${OPENAI_API_KEY}
      chat:
        options:
          model: gpt-4o-mini
          temperature: 0.0

A high temperature makes the model more creative, so the same user query can be transformed differently in each call. For that reason, the Spring AI documentation recommends using a low temperature such as 0.0 to get more deterministic results.

4. Modular RAG

In Spring AI, RetrievalAugmentationAdvisor ties the RAG stages together. Once we register it on a ChatClient, every prompt runs through the full pipeline, e.g.:

RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
  .queryTransformers(
    RewriteQueryTransformer.builder()
      .chatClientBuilder(chatClientBuilder)
      .build(),
    TranslationQueryTransformer.builder()
      .chatClientBuilder(chatClientBuilder)
      .targetLanguage("english")
      .build(),
    CompressionQueryTransformer.builder()
      .chatClientBuilder(chatClientBuilder)
      .build())
  .documentRetriever(VectorStoreDocumentRetriever.builder()
    .vectorStore(vectorStore)
    .topK(3)
    .build())
  .build();

ChatClient chatClient = chatClientBuilder.defaultAdvisors(ragAdvisor)
  .build();

Here, vectorStore is a VectorStore instance holding our embedded documents, e.g. backed by an Oracle vector database.

The builder of the RetrievalAugmentationAdvisor exposes at least one method for each RAG stage. In our example, we configure queryTransformers() for the pre-retrieval stage and the documentRetriever() method for the retrieval stage. Let’s focus on the former.

The queryTransformers() method accepts multiple query transformers. They implement the QueryTransformer interface and run in sequence, each one receiving the previous one’s output. Each transformer makes its own LLM call that requires a ChatClient.Builder.

Spring AI provides three query transformers. To demonstrate each transformer separately, we’ll implement a test method for each in our test class:

@SpringBootTest
class PreRetrievalQueryTransformerLiveTest {
    private static final Logger logger = LoggerFactory.getLogger(PreRetrievalQueryTransformerLiveTest.class);

    @Autowired
    private ChatClient.Builder chatClientBuilder;
}

Let’s check them out.

5. RewriteQueryTransformer

RewriteQueryTransformer filters out irrelevant information from a query. By keeping only the relevant parts, it produces a concise text that we expect to find better results. This transformer is useful when the user query is lengthy and unclear.

As an example, let’s show how it handles a verbose query:

@Test
void whenQueryIsVerbose_thenRewrite() {
    QueryTransformer queryTransformer = RewriteQueryTransformer.builder()
      .chatClientBuilder(chatClientBuilder)
      .build();

    String userQuery = "I'm studying machine learning and honestly the whole thing is confusing. What is an LLM anyway?";
    Query transformed = queryTransformer.transform(new Query(userQuery));

    logger.info("Rewrite: [{}] -> [{}]", userQuery, transformed.text());
}

We can see the following or similar log when executing the test:

Rewrite: [I'm studying machine learning and honestly the whole thing is confusing. What is an LLM anyway?] -> [What is an LLM in machine learning?]

The transformer identifies machine learning and LLM as important context to keep and discards the user’s introduction that is irrelevant for the search.

6. TranslationQueryTransformer

Embedding models work best for the languages they were trained on. For example, if most of the training data is in English, the model represents English texts more precisely than other written languages. Consequently, a query written in another language such as Chinese or French may align poorly with documents stored in English, even when they have the same meaning.

TranslationQueryTransformer solves this by translating the user’s query into the target language: the language of the documents in our datastore.

To see this, let’s add a new test method to send a Traditional Chinese query 什麼是向量資料庫？ asking “What is a vector database?” to the transformer:

@Test
void whenQueryIsInTraditionalChinese_thenTranslate() {
    QueryTransformer queryTransformer = TranslationQueryTransformer.builder()
      .chatClientBuilder(chatClientBuilder)
      .targetLanguage("english")
      .build();

    String userQuery = "什麼是向量資料庫？";
    Query transformed = queryTransformer.transform(new Query(userQuery));

    logger.info("Translation: [{}] -> [{}]", userQuery, transformed.text());
}

We specify the target language explicitly using the targetLanguage() method. The log shows the Chinese query has been translated correctly:

Translation: [什麼是向量資料庫？] -> [What is a vector database?]

If the query is already in the target language or in a language the LLM can’t recognize, it’s returned unchanged. For example, if we use a query in English, the output is identical:

Translation: [What is a vector database?] -> [What is a vector database?]

Similarly, a query with made-up words of an unknown language passes through unchanged as well:

Translation: [Gblorvak tensig quimbo darnel?] -> [Gblorvak tensig quimbo darnel?]

7. CompressionQueryTransformer

A user query usually refers to previous messages and doesn’t contain complete information on its own. For that reason, a single query in a chat may be ambiguous and not return anything useful from the vector store.

CompressionQueryTransformer solves this by compressing the conversation history and using it to rewrite the query to resolve possible ambiguities.

Unlike the previous two transformers, it requires the history of the Query, so must supply both the query text and the history:

@Test
void whenQueryIsFollowUpWithHistory_thenCompress() {
    QueryTransformer queryTransformer = CompressionQueryTransformer.builder()
      .chatClientBuilder(chatClientBuilder)
      .build();

    List<Message> history = List.of(
      new UserMessage("What is Spring AI?"),
      new AssistantMessage("Spring AI is an application framework that brings Spring's design principles to AI engineering."));
    
    String followUpQuery = "And how does it support RAG?";
    
    Query query = Query.builder()
      .text(followUpQuery)
      .history(history)
      .build();

    Query transformed = queryTransformer.transform(query);

    logger.info("Compression: [{}] -> [{}]", followUpQuery, transformed.text());
}

Here, we construct the conversation history explicitly as a list of messages. Executing the test should produce a log like this one:

Compression: [And how does it support RAG?] -> [How does Spring AI support Retrieval-Augmented Generation (RAG)?]

The original query was unclear on its own because of “it.” The transformer compresses it and the conversation history into a new query where the meaning of “it” is clear.

In a real application, a ChatClient doesn’t keep previous messages by default, so we need to register a chat memory advisor with the RAG advisor. Otherwise, the transformer receives no history to work with.

However, the chat history may not always be relevant as the user may switch the topic during conversation. Let’s run the same test again with an unrelated query:

String followUpQuery = "What's a good recipe for carbonara?";

The transformer should recognize the topic switch and not include anything from previous messages:

Compression: [What's a good recipe for carbonara?] -> [Can you provide a good recipe for carbonara?]

8. Query Transformer Considerations

All query transformers have these common characteristics:

• They run sequentially, and each receives the previous one’s output.
• Each transformer adds an LLM call.
• Temperature needs to be low (e.g., 0.0) to ensure more deterministic results.

Because of additional LLM calls, the overall latency adds up. So, we should enable only the transformers we actually need:

• RewriteQueryTransformer simplifies verbose and unclear queries to provide clarity.
• TranslationQueryTransformer translates queries into the language of the database.
• CompressionQueryTransformer compresses and merges a query with its history into a new query containing relevant context from the chat history.

9. Conclusion

In this article, we explained query transformation in the pre-retrieval stage of Spring AI’s modular RAG architecture.

Three QueryTransformer implementations that Spring AI provide transform queries to improve retrieval performance.

As always, the full source code is available over on GitHub.