Optimizing Search Queries with LLMs

Welcome to the first lesson in our "Advanced RAG" module. In our previous lessons, you've built a robust retrieval system, culminating in a hybrid search engine that combines the power of semantic and keyword search. You now have a system that can find relevant documents with impressive accuracy. However, as with any system that relies on user input, its performance is often limited by the quality of that input. A vague, ambiguous, or poorly phrased query can lead to poor retrieval, no matter how sophisticated your search algorithm is.

Today, we'll tackle this "garbage in, garbage out" problem head-on. Our goal is to improve retrieval relevance by transforming user queries into optimized search queries using an LLM. Instead of just passively accepting the user's query, we will actively use another LLM to refine, rephrase, and expand it before it ever hits our vector database. This is a crucial step in building production-grade AI systems that feel truly intelligent and helpful.

The "Lost in Translation" Problem

The core issue is that users speak in natural, everyday language, while your knowledge base might be written in a more formal, technical, or domain-specific style. This creates a "vocabulary gap" that can prevent a successful match even when the right information exists.

As the article highlights, the goal of query transformation is to act as a translator between "human-speak" and "document-speak." We use an LLM to anticipate how the information might be stored and rewrite the user's query to better align with the vector space of our documents.

This transformation step happens right at the beginning of the retrieval process, as shown in the architectural diagram below.

Technique 1: Multi-Query Generation

One of the most effective and intuitive transformation techniques is Multi-Query Generation. Instead of trying to find the single "perfect" version of a query, we use an LLM to generate several different versions from multiple perspectives. This increases the chances of finding relevant documents that might have been missed by the original phrasing.

The intuition is simple: by casting a wider net, we are more likely to catch the documents we need. A slight rephrasing might be all it takes to move the query's embedding closer to a relevant document chunk in the high-dimensional vector space.

The process for a Multi-Query RAG system is as follows:

1. Generate: Take the original user query and pass it to an LLM with a prompt asking it to generate N different versions of the question.
2. Retrieve: For each of the N generated queries, perform a standard vector search against your database to retrieve a list of relevant documents.
3. Combine: Merge all the retrieved document lists into a single list, removing any duplicates.
4. Synthesize: Pass the final, unique set of documents as context to another LLM call to synthesize the final answer based on the original question.

Hands-On: Implementing Multi-Query Retrieval

Let's integrate this into your RAG system. We'll write a function that uses the OpenAI API to generate query variations and then modify our search logic to use them. This implementation avoids external frameworks to ensure you understand the core mechanics.

The following script builds on the code from our previous lessons.

import chromadb
import os
from openai import OpenAI
import pprint
import json

# --- Setup ---
client = OpenAI() # Assumes OPENAI_API_KEY is set
chroma_client = chromadb.PersistentClient(path="./my_vector_db")
collection_name = "rag_system_docs"
collection = chroma_client.get_collection(name=collection_name)

# --- 1. Query Transformation Step ---
def generate_multiple_queries(query: str, num_queries: int = 4) -> list[str]:
    """
    Generates multiple versions of a query using an LLM.
    """
    prompt = f"""
    You are an AI language model assistant. Your task is to generate {num_queries}
    different versions of the given user question to retrieve relevant documents from a vector
    database. By generating multiple perspectives on the user question, your goal is to help
    the user overcome some of the limitations of distance-based similarity search.

    Provide these alternative questions as a JSON list of strings.
    
    Original question: {query}
    """
    
    print("--- Generating multiple query perspectives ---")
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant that generates multiple search queries based on a single input query."},
            {"role": "user", "content": prompt}
        ],
        response_format={"type": "json_object"}
    )
    
    try:
        query_list = json.loads(response.choices[0].message.content)
        # Assuming the JSON response is a dict with a key like "queries"
        # This might need adjustment based on consistent LLM output format
        if isinstance(query_list, dict):
            # A common pattern is {"queries": ["q1", "q2", ...]}
            actual_list = next(iter(query_list.values()))
            if isinstance(actual_list, list):
                 print("Generated queries:", actual_list)
                 return actual_list
        raise ValueError("JSON from LLM is not in the expected list format")
    except (json.JSONDecodeError, ValueError, TypeError) as e:
        print(f"Error parsing LLM response: {e}")
        print(f"Raw response: {response.choices[0].message.content}")
        # Fallback to using the original query
        return [query]

# --- 2. Retrieval and Combination Step ---
def retrieve_and_combine(queries: list[str], n_results_per_query: int = 3):
    """
    Retrieves documents for each query and returns a unique, combined set.
    """
    all_retrieved_docs = {} # Using a dict to handle duplicates easily
    
    for q in queries:
        print(f"--- Retrieving for query: '{q}' ---")
        results = collection.query(
            query_texts=[q],
            n_results=n_results_per_query
        )
        
        # Add documents to our dict, using ID as the key to automatically handle duplicates
        for i, doc_id in enumerate(results['ids'][0]):
            if doc_id not in all_retrieved_docs:
                all_retrieved_docs[doc_id] = {
                    "content": results['documents'][0][i],
                    "metadata": results['metadatas'][0][i]
                }
    
    # The values of the dictionary are the unique documents
    unique_docs = list(all_retrieved_docs.values())
    print(f"\nRetrieved {len(unique_docs)} unique documents in total.")
    return unique_docs

# --- 3. Synthesis Step ---
def synthesize_answer(query: str, context_docs: list[dict]):
    """
    Generates a final answer based on the original query and retrieved context.
    """
    context_str = "\n\n---\n\n".join([doc['content'] for doc in context_docs])
    
    prompt = f"""
    Based on the following context, please provide a comprehensive answer to the user's original question.
    
    Context:
    {context_str}
    
    Original Question: {query}
    
    Answer:
    """
    
    print("\n--- Synthesizing final answer ---")
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant that answers questions based on provided context."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content

# --- Run the full pipeline ---
user_query = "What is task decomposition for LLM agents?"

# 1. Generate queries
generated_queries = generate_multiple_queries(user_query)

# 2. Retrieve and combine
retrieved_context = retrieve_and_combine(generated_queries)

# 3. Synthesize answer
final_answer = synthesize_answer(user_query, retrieved_context)

print("\n\n--- FINAL ANSWER ---")
print(final_answer)

From Multi-Query to RAG-Fusion

The simple Multi-Query approach of taking a unique union of documents is a great start, but it treats all retrieved documents equally. What if one document was a top-ranked result for three of our generated queries, while another was a low-ranked result for only one? Intuitively, the first document is more likely to be relevant.

This is where RAG-Fusion comes in. It improves upon Multi-Query by using the Reciprocal Rank Fusion (RRF) algorithm you learned in the last lesson. Instead of just de-duplicating, we use RRF to re-rank the combined results, giving a score boost to documents that appear consistently across the different query result lists.

This beautifully connects the previous lesson with this one, showing how these advanced techniques can be layered to build an even more powerful system.

System Design and Trade-Offs

As a future AI Engineer, it's crucial to think about the practical implications of these techniques. Adding query transformation isn't "free." It introduces trade-offs in terms of cost, latency, and complexity.

• Cost: Each query transformation is an extra LLM call, which adds to your API bill. A Multi-Query approach with 4 generated queries means you are now performing 4x the number of vector searches.
• Latency: The additional LLM call and multiple retrieval steps add time to the overall response generation process.
• Complexity: The logic becomes more intricate. You have to manage the transformation step, the multi-retrieval loop, and the fusion logic.

The key is to decide when this added overhead is justified by the improvement in retrieval quality. For many production use cases, the answer is a firm "yes," as it dramatically improves the user experience and reduces instances of the RAG system failing to find an answer.

Here's a useful heuristic for when to use different query transformation techniques:

Conclusion

You have now taken a major step into advanced RAG by learning to actively transform user queries. You've learned that a user's input isn't a final command but a starting point that can be optimized with an LLM. By implementing the Multi-Query technique, you've made your RAG system more resilient to ambiguous or poorly phrased questions. Furthermore, by understanding how to combine this with RRF to create RAG-Fusion, you're connecting concepts to build progressively more sophisticated systems.

Key Takeaways:

• Query Transformation is a critical step in advanced RAG to bridge the "vocabulary gap" between users and documents.
• Multi-Query Generation uses an LLM to create multiple perspectives of a user's question, increasing the chances of retrieving relevant context.
• RAG-Fusion combines Multi-Query with Reciprocal Rank Fusion (RRF) to intelligently re-rank documents based on their prevalence across multiple searches.
• These techniques add cost and latency, but often provide a worthwhile improvement in retrieval quality for production systems.

In our next lesson, we will explore another powerful query transformation technique: using an LLM to improve retrieval coverage by expanding user queries with generated sub-questions. This is particularly useful for complex questions that can't be answered by a single piece of context.