Enhancing Retrieval with Sub-Question Generation

Hello! In our last lesson, we significantly enhanced your RAG system by implementing multi-query generation and RAG-Fusion. You learned to transform a user's query into multiple, differently-phrased versions to improve the relevance of retrieved documents. This was about tackling the "vocabulary gap" and finding the right documents even if the user's wording didn't perfectly match your knowledge base.

Today, we'll address a different, but equally important, challenge: complex questions. Often, users don't ask simple, single-fact questions. They ask for comparisons, summaries of multiple topics, or cause-and-effect analyses. Answering these queries requires gathering information from several different documents or contexts. Our goal for this lesson is to improve retrieval coverage by expanding user queries with generated sub-questions. This technique, known as Query Decomposition, is a cornerstone of building advanced, capable RAG systems.

From Rephrasing to Decomposing

The key difference between our previous lesson and this one is the shift from rephrasing to decomposing.

• Multi-Query (Rephrasing): Takes one question and generates multiple ways of asking the same thing. This improves the chance of a semantic match.
• Query Decomposition: Takes one complex question and breaks it into several simpler, independent sub-questions. This ensures we gather all the distinct pieces of information needed to construct a comprehensive answer.

Consider the query: "Who has more siblings, Jamie or Sansa from Game of Thrones, and what are their respective family names?" A single vector search is highly unlikely to find a document that perfectly answers this comparative and multi-part question. Query decomposition allows us to approach this systematically.

The core idea is to use an LLM as a "planner." It analyzes the user's request and creates a "to-do list" of simpler questions. We then execute this to-do list by running a separate retrieval for each item.

The Query Decomposition Workflow

The process is a classic "Map-Reduce" pattern, which should be familiar from your software engineering background.

1. Decompose (Map): A single complex query is "mapped" to multiple sub-questions by an LLM.
2. Retrieve (Map): Each sub-question is used to retrieve relevant documents from the vector database. This step is highly parallelizable, which is a key advantage for system performance.
3. Synthesize (Reduce): All the retrieved documents are gathered and "reduced" along with the original query into a final, synthesized answer by another LLM call.

This workflow is illustrated perfectly in the diagram below. A single, complex question is broken into Q1, Q2, and Q3. Each sub-question has its own independent retrieval path, and the results are combined in the end to form the final answer.

Tying it to Production: Parallel Tool Use

This pattern of breaking a task into parallel sub-tasks is not just an academic exercise; it's a core feature of modern agentic frameworks and LLM APIs. For instance, models with "tool calling" or "function calling" capabilities can perform this decomposition naturally.

When you give a capable LLM a set of tools (like a search_docs function) and a complex query, it can decide to call the same tool multiple times with different parameters in parallel. This is Query Decomposition in action.

Hands-On: Implementing Query Decomposition

Now, let's implement this yourself. We will add a decompose_query function to our RAG pipeline. The other functions for retrieval and synthesis can be reused from our previous lesson, highlighting the modularity of this architecture.

Here is the Python code. We'll focus on the new decomposition function and how it slots into the existing pipeline.

import chromadb
import os
from openai import OpenAI
import json




# --- Setup (Same as previous lesson) ---
client = OpenAI() # Assumes OPENAI_API_KEY is set
chroma_client = chromadb.PersistentClient(path="./my_vector_db")
collection_name = "rag_system_docs"



# Make sure your collection exists before running.
# collection = chroma_client.create_collection(name=collection_name)
collection = chroma_client.get_collection(name=collection_name)




# --- 1. NEW: Query Decomposition Step ---
def decompose_query_into_subquestions(query: str) -> list[str]:
    """
    Decomposes a complex query into a list of simpler sub-questions using an LLM.
    """
    prompt = f"""
    You are an expert at breaking down complex user questions into a series of simpler, self-contained sub-questions.
    Your goal is to generate a list of questions that can be answered independently by a search engine to gather all the necessary information to fully answer the original complex question.

    - Each sub-question should be a complete, standalone question.
    - Do not lose any part of the original question's intent.
    - If the question is already simple and cannot be broken down, return it as a single-element list.

    Provide the output as a JSON object with a single key "sub_questions" which contains a list of strings.

    Original complex question: "{query}"
    """
    
    print("--- Decomposing query into sub-questions ---")
    try:
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are a helpful assistant that decomposes complex questions into sub-questions."},
                {"role": "user", "content": prompt}
            ],
            response_format={"type": "json_object"}
        )
        
        response_dict = json.loads(response.choices[0].message.content)
        sub_questions = response_dict.get("sub_questions", [query])
        



        # Basic validation to ensure we got a list of strings
        if isinstance(sub_questions, list) and all(isinstance(q, str) for q in sub_questions):
            print("Generated sub-questions:", sub_questions)
            return sub_questions
        else:
            print("Warning: LLM did not return a valid list of strings. Falling back to original query.")
            return [query]

    except (json.JSONDecodeError, KeyError, TypeError) as e:
        print(f"Error during decomposition: {e}. Falling back to original query.")
        return [query]




# --- 2. REUSED: Retrieval and Combination Step ---
def retrieve_and_combine(queries: list[str], n_results_per_query: int = 2):
    """
    Retrieves documents for each query and returns a unique, combined set.
    This function is reused from our previous lesson.
    """
    all_retrieved_docs = {}
    
    for q in queries:
        print(f"--- Retrieving for sub-question: '{q}' ---")



        # Assuming your collection is not empty and has been indexed
        results = collection.query(
            query_texts=[q],
            n_results=n_results_per_query
        )
        
        if results['ids'][0]:
            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]
                    }
    
    unique_docs = list(all_retrieved_docs.values())
    print(f"\nRetrieved {len(unique_docs)} unique documents in total from all sub-questions.")
    return unique_docs




# --- 3. REUSED: Synthesis Step ---
def synthesize_answer(query: str, context_docs: list[dict]):
    """
    Generates a final answer based on the original query and retrieved context.
    This function is also reused from our previous lesson.
    """
    if not context_docs:
        return "I could not find any relevant information to answer your question."
    
    context_str = "\n\n---\n\n".join([doc['content'] for doc in context_docs])
    
    prompt = f"""
    You are a helpful AI assistant. Your task is to provide a comprehensive answer to the user's original question based on the provided context.
    The context is a collection of documents retrieved in response to several sub-questions derived from the original question.
    Synthesize the information from these documents to form a coherent answer.

    Context:
    {context_str}
    
    Original Question: {query}
    
    Answer:
    """
    
    print("\n--- Synthesizing final answer from combined context ---")
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are an expert Q&A assistant."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content




# --- Run the full pipeline ---
# Make sure to populate your ChromaDB collection with some documents first!
# For example, documents about different smart phones.
user_query = "Compare the battery life and camera quality of the iPhone 15 Pro and the Google Pixel 8 Pro, and also mention which one has a better display."




# 1. Decompose the query
sub_questions = decompose_query_into_subquestions(user_query)




# 2. Retrieve context for all sub-questions
retrieved_context = retrieve_and_combine(sub_questions)




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

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

System Design and Interview Prep

For your goal of becoming an AI Engineer at a top tech company, understanding the "why" is as important as the "how."

• When to Decompose: This technique is most valuable for complex queries that involve comparison, summarization of multiple points, or sequential reasoning. A simple heuristic could be to run a preliminary check with an LLM: "Is this question complex and multi-part? True or False." If true, trigger the decomposition flow.
• Trade-offs: Decomposition adds latency and cost (one extra LLM call upfront, multiple database lookups). The benefit is a massive improvement in coverage and the ability to answer questions that are impossible for a simpler RAG system. The retrieval step is parallelizable, so with the right infrastructure (like microservices or async workers), the added latency can be managed.
• Combining Techniques: You can use decomposition and multi-query rephrasing together. First, decompose the complex query into sub-questions. Then, for each sub-question, you could generate multiple rephrasings before retrieval. This creates a very thorough but also very expensive and high-latency pipeline, reserved only for the most critical and complex queries.

Conclusion

You have now learned a second powerful query transformation technique. While multi-query rephrasing helps you find better matches for a single topic, query decomposition helps you cover all the topics embedded within a complex user question. By using an LLM to break down a query into its constituent parts, you can ensure your RAG system retrieves all the necessary context to deliver a complete and comprehensive answer.

Key Takeaways:

• Query Decomposition breaks complex user questions into simpler, independent sub-questions.
• This technique improves retrieval coverage for multi-faceted queries.
• The workflow follows a Map-Reduce pattern: decompose, retrieve in parallel for each sub-question, and synthesize a final answer.
• Modern agentic systems often implement this pattern through parallel tool calling.
• There is a clear trade-off between the increased cost/latency and the significant improvement in answer completeness.

In our previous lessons, we've focused on getting more and better documents into our context. But what happens when we have a lot of documents, some more relevant than others? In our next lesson, we will address this by learning how to apply a re-ranking model to re-order retrieved contexts for maximum relevance, ensuring the final LLM sees the most important information first.