Capsule 05: Query decomposition — when one question is several questions

Capsule overview

So far we've seen techniques for queries that are badly expressed: ambiguous (expansion), incomplete (rewriting with context), keyword-style (structural rewriting). This capsule covers a different case: queries that are well expressed but contain several questions in one.

"Compare FastAPI vs Flask for authentication, and recommend which one to use in production." "How do I set up PostgreSQL in Docker, connect FastAPI with SQLAlchemy, and deploy to AWS?" "Why is HNSW faster than IVF, and when is it worth using PQ?"

These queries are perfectly clear. The user knows what they want. The problem is that the answer requires chunks from several different documents. A single search with the embedding of the full query returns docs that touch on some aspects but rarely cover all of them. Decomposition splits the query into independent sub-queries, searches each one separately, and combines the results.

This capsule teaches you to detect when a query is a candidate for decomposition (not all are), how to decompose correctly while preserving the original question, and how to aggregate the results without flooding the LLM's context.

By the end of this capsule you'll be able to:

• ✅ Identify the three classes of query that benefit from decomposition: comparative, multi-step, causal
• ✅ Implement decomposition with an LLM and structured outputs
• ✅ Search sub-queries in parallel and aggregate results with deduplication
• ✅ Decide when decomposition adds value vs when the direct query is enough
• ✅ Anticipate the "context explosion" trap: dumping 30 chunks floods the LLM
• ✅ Combine decomposition with retrieve-then-rerank for maximum precision

Estimated time: 30-35 minutes

When to use decomposition

Three query patterns where decomposition is the right tool:

Pattern 1: comparative queries

"Compare FastAPI vs Flask for authentication"

→ Sub-query 1: "How does authentication work in FastAPI?"
→ Sub-query 2: "How does authentication work in Flask?"
→ Sub-query 3: "What are the key differences between FastAPI and Flask in auth?"

Without decomposition, a single search with the full query will probably return docs about one of the two frameworks but not the other — the query "compare A vs B" has an embedding that doesn't point clearly at A or at B.

Pattern 2: multi-step queries

"How do I set up PostgreSQL in Docker, connect FastAPI with SQLAlchemy, and deploy to AWS?"

→ Sub-query 1: "How do I configure PostgreSQL in a Docker container?"
→ Sub-query 2: "How do I connect FastAPI with SQLAlchemy to a PostgreSQL database?"
→ Sub-query 3: "How do I deploy a FastAPI application to AWS?"

Each sub-query maps to an independent aspect of the problem. A single search retrieves docs that touch every topic tangentially but go deep on none.

Pattern 3: causal queries with a follow-up

"Why is HNSW faster than IVF, and when is it worth using PQ?"

→ Sub-query 1: "Why is HNSW faster than IVF? Internal mechanism"
→ Sub-query 2: "When do you use PQ instead of HNSW or IVF?"
→ Sub-query 3 (optional): "Trade-offs between HNSW, IVF and PQ"

Causal queries tend to mix "why" with "how to apply" — those are two different questions that need two different answers.

When NOT to decompose

Simple heuristic: if the query uses words like "compare", "differences between", "and also", "plus", or has several interrogative sentences, it's a candidate. If it's a single direct question, it isn't.

Implementation with structured outputs

# query_decomposition.py
from openai import OpenAI
from pydantic import BaseModel, Field
from typing import List
import os


client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


class DecomposedQuery(BaseModel):
    is_decomposable: bool = Field(
        description="True if the query benefits from decomposition (multi-aspect, comparative, multi-step). False if it's a single direct question."
    )
    sub_queries: List[str] = Field(
        description="If decomposable, list of independent sub-queries that together cover the original. Empty if not decomposable.",
        max_items=6,
    )
    reasoning: str = Field(
        description="One-sentence explanation of why the query is or isn't decomposable"
    )


SYSTEM_PROMPT = """You are an expert at analyzing search queries and decomposing complex ones.

Given a user query, decide:
1. If it's a SIMPLE single question → mark as not decomposable, return empty sub-queries
2. If it's COMPARATIVE (compares 2+ things) → decompose into individual aspect queries + comparison
3. If it's MULTI-STEP (asks for sequential steps) → decompose into one query per step
4. If it's CAUSAL+APPLICATION (mixes "why" with "how") → decompose into separate why/how queries

Rules:
- Each sub-query must be ANSWERABLE INDEPENDENTLY
- Together, sub-queries must cover the FULL original query
- Use natural language for sub-queries
- Maximum 6 sub-queries (typically 3-4 is ideal)

Examples:

Query: "How do I implement OAuth2 in FastAPI?"
→ NOT decomposable. Single direct question.

Query: "Compare FastAPI and Flask for authentication"
→ DECOMPOSABLE:
   1. "How does authentication work in FastAPI?"
   2. "How does authentication work in Flask?"
   3. "What are the key differences between FastAPI and Flask authentication?"

Query: "How do I deploy a FastAPI app with PostgreSQL on AWS?"
→ DECOMPOSABLE:
   1. "How do I configure PostgreSQL for production?"
   2. "How do I connect FastAPI to PostgreSQL using SQLAlchemy?"
   3. "How do I deploy a FastAPI application to AWS?"
"""


def decompose_query(query: str) -> DecomposedQuery:
    """Analyze and possibly decompose a complex query."""
    response = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f'Query: "{query}"\n\nAnalyze and decompose if appropriate.'},
        ],
        response_format=DecomposedQuery,
        temperature=0.2,
    )
    return response.choices[0].message.parsed


# Try it out
queries = [
    "How do I implement OAuth2 in FastAPI?",
    "Compare FastAPI and Flask for authentication",
    "How do I deploy a FastAPI app with PostgreSQL on AWS?",
    "What is HNSW?",
    "Why is HNSW faster than IVF, and when should I use PQ?",
]

for q in queries:
    result = decompose_query(q)
    print(f"\nQuery: {q}")
    print(f"  Decomposable: {result.is_decomposable}")
    print(f"  Reasoning: {result.reasoning}")
    if result.sub_queries:
        print(f"  Sub-queries:")
        for i, sq in enumerate(result.sub_queries, 1):
            print(f"    {i}. {sq}")

Expected output:

Query: How do I implement OAuth2 in FastAPI?
  Decomposable: False
  Reasoning: Single direct question with one specific topic, no comparison or multi-step structure.

Query: Compare FastAPI and Flask for authentication
  Decomposable: True
  Reasoning: Comparative query requiring information about each framework separately and their differences.
  Sub-queries:
    1. How does authentication work in FastAPI?
    2. How does authentication work in Flask?
    3. What are the key differences between FastAPI and Flask authentication?

Query: How do I deploy a FastAPI app with PostgreSQL on AWS?
  Decomposable: True
  Reasoning: Multi-step query covering three independent topics: PostgreSQL setup, FastAPI integration, and AWS deployment.
  Sub-queries:
    1. How do I configure PostgreSQL for production use?
    2. How do I connect a FastAPI application to PostgreSQL using SQLAlchemy?
    3. How do I deploy a FastAPI application to AWS?

Query: What is HNSW?
  Decomposable: False
  Reasoning: Definitional question with one direct answer.

Query: Why is HNSW faster than IVF, and when should I use PQ?
  Decomposable: True
  Reasoning: Causal-and-application query mixing "why" comparison with "when to use" application advice.
  Sub-queries:
    1. Why is HNSW faster than IVF for vector search?
    2. When should I use PQ instead of HNSW or IVF?

The advantage of the structured output with is_decomposable: the model decides whether to decompose or not. You don't have to call the LLM, parse it, and then check whether the decomposition makes sense. A single call settles both things.

Full pipeline: decompose → search in parallel → aggregate

# decomposition_pipeline.py
from concurrent.futures import ThreadPoolExecutor
from collections import OrderedDict


def search_with_decomposition(query: str, top_k_per_subquery: int = 5):
    """
    Full pipeline:
    1. Decompose if it applies
    2. Search each sub-query in parallel
    3. Aggregate deduplicated results
    """
    decomp = decompose_query(query)

    # Case 1: simple query, don't decompose
    if not decomp.is_decomposable:
        results = collection.query(query_texts=[query], n_results=top_k_per_subquery * 2)
        return {
            "is_decomposed": False,
            "sub_queries": [query],
            "documents": results['documents'][0],
            "metadatas": results['metadatas'][0],
            "ids": results['ids'][0],
        }

    # Case 2: complex query, decompose
    queries_to_search = decomp.sub_queries

    def search_one(q):
        return collection.query(query_texts=[q], n_results=top_k_per_subquery)

    # Parallel: 3-5 sub-queries in parallel is trivial
    with ThreadPoolExecutor(max_workers=5) as executor:
        all_results = list(executor.map(search_one, queries_to_search))

    # Aggregate results with deduplication (keeping the first appearance)
    unique_results = OrderedDict()
    for sub_query_results in all_results:
        for doc_id, doc, meta in zip(
            sub_query_results['ids'][0],
            sub_query_results['documents'][0],
            sub_query_results['metadatas'][0],
        ):
            if doc_id not in unique_results:
                unique_results[doc_id] = (doc, meta)

    return {
        "is_decomposed": True,
        "sub_queries": queries_to_search,
        "documents": [d for d, m in unique_results.values()],
        "metadatas": [m for d, m in unique_results.values()],
        "ids": list(unique_results.keys()),
    }

Limiting the context sent to the LLM

If you decompose into 4 sub-queries and each one returns 5 docs, you end up with ~20 unique docs to pass to the LLM. That's too much context — the LLM gets distracted by noise and answer quality drops ("lost in the middle").

Fix: rerank after aggregating.

def decompose_then_rerank(query: str, final_top_k: int = 8):
    """
    Decompose + aggregate + rerank to limit the final context sent to the LLM.
    """
    aggregated = search_with_decomposition(query, top_k_per_subquery=5)

    # If it didn't decompose, return directly
    if not aggregated["is_decomposed"]:
        return aggregated["documents"][:final_top_k]

    # Rerank with the ORIGINAL query (not the sub-queries)
    # to favor docs that address the whole, not isolated aspects
    reranked = cross_encoder_rerank(
        query=query,
        documents=aggregated["documents"],
        top_k=final_top_k,
    )

    return reranked

Why rerank with the original query (not the sub-queries): the LLM's answer has to address the user's full query. Documents that only cover one aspect lose relevance against documents that cover several aspects. The cross-encoder with the original query favors holistic docs.

Generating the final answer

When you decompose, the LLM has to synthesize information from several sub-topics. The prompt needs to be explicit:

def generate_decomposed_answer(query: str, retrieved_docs: list[str]) -> str:
    """Generate the final answer addressing all aspects of the original query."""
    context = "\n\n---\n\n".join(retrieved_docs)

    prompt = f"""Answer the user's question comprehensively. The question may have
multiple aspects — make sure to address ALL of them in a structured response.

Use ONLY the provided context. If a specific aspect isn't covered in the context,
say so explicitly rather than making up information.

Question: {query}

Context:
{context}

Answer (address all aspects):"""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "You are a technical assistant that provides structured, comprehensive answers."},
            {"role": "user", "content": prompt},
        ],
        temperature=0.0,
    )
    return response.choices[0].message.content

The secret: the instruction "address ALL of them" + "if a specific aspect isn't covered, say so explicitly" prevents partial answers. If the decomposition retrieved docs about 2 of 3 sub-topics, the LLM will say "I found no information about the third aspect" instead of making something up.

Traps and common mistakes

Trap 1: decomposing simple queries

The mistake: decomposition turned on for every query.

Symptom: a simple query like "what is HNSW?" gets decomposed into "HNSW definition" + "what HNSW is for" + "how HNSW works". The system runs 3 unnecessary searches.

How to prevent it: the structured output with is_decomposable=False prevents it if the prompt is well designed. Validate with an eval set.

Trap 2: context explosion

The mistake: you decompose into 5 sub-queries × 10 docs each = 50 unique docs to the LLM.

Symptom: latency spikes, answer quality drops ("lost in the middle"), generation cost multiplies.

How to prevent it: after aggregating, rerank and cap at 5-10 final docs.

Trap 3: sub-queries that aren't independent

The mistake: sub-queries that depend on the result of the previous ones.

Query: "Which is the best Python framework for my case?"
Bad sub-queries:
  1. "What are the popular Python frameworks?"
  2. "Which of the above is best?"  ← depends on 1, can't be searched alone

Symptom: searches for sub-queries that depend on each other come back with poor results.

How to prevent it: an explicit prompt saying "each sub-query must be ANSWERABLE INDEPENDENTLY". If you can't phrase independent sub-queries, don't decompose.

Trap 4: naive deduplication

The mistake:

unique_docs = list(set(all_documents))  # set by doc content

Symptom: near-identical docs (with minor whitespace changes) count as different. The context fills up with duplicates.

How to prevent it: deduplicate by doc_id, not by content. IDs are stable.

Trap 5: using the sub-queries to rerank instead of the original query

The mistake: rerank with each sub-query, then aggregate the rankings.

Symptom: the rerank favors docs that answer individual aspects well instead of docs that cover the whole question.

How to prevent it: rerank with the user's original query. The sub-queries are only for the initial retrieval.

Trap 6: decomposition with no fallback when a sub-query fails

The mistake: one of the sub-queries returns zero results (topic not in the corpus). You treat that as zero contribution.

Symptom: the LLM's answer ignores that aspect without warning anyone.

How to prevent it: detect sub-queries with zero results and tell the LLM:

sub_query_coverage = {}
for sub_q, results in zip(sub_queries, all_results):
    sub_query_coverage[sub_q] = len(results['documents'][0])

# In the prompt to the LLM:
missing_aspects = [sq for sq, count in sub_query_coverage.items() if count == 0]
if missing_aspects:
    prompt += f"\n\nNote: No information was found for: {missing_aspects}. Mention this explicitly in your answer."

Applied exercise

Scenario: you're the AI Engineer at a programming education platform. The RAG system serves students asking questions about backend courses.

Query log analysis:

Category                            %    Example
──────────────────────────────────────────────────────────────────
Simple direct question             55%   "what is REST?"
Comparison between frameworks      18%   "Django vs FastAPI vs Flask"
Multi-step (build + deploy)        15%   "how to build and deploy an API"
Causal + application                8%   "why is async faster and how do I use it?"
Vague/incomplete                    4%   "help with my project"

Current metrics (no decomposition):

• Precision@5: 81%
• For simple queries: 88%
• For comparative queries: 62% ← problem
• For multi-step queries: 68% ← problem

Your job:

1. Decide whether decomposition is worth adding.
2. Design the pipeline including a dynamic skip.
3. Estimate the impact on latency and cost.

Category                   %     Pre-decomp     Post-decomp     Weighted gain
──────────────────────────────────────────────────────────────────────────
Simple                   55%       88%            88% (unchanged)         0
Comparative              18%       62%            85% (estimated)         +4.1 pts
Multi-step               15%       68%            85% (estimated)         +2.6 pts
Causal+application        8%       65% (assumed)  82% (estimated)         +1.4 pts
Vague                     4%       50%            50% (N/A)               0

Total expected gain in precision@5: +8 points (81% → 89%)

def smart_pipeline(query: str, final_top_k: int = 5) -> list[str]:
    """
    Pipeline with conditional decomposition.
    """
    # The LLM decides whether to decompose (structured output with is_decomposable)
    decomp = decompose_query(query)

    if not decomp.is_decomposable:
        # Standard pipeline: retrieve + rerank
        results = collection.query(query_texts=[query], n_results=20)
        return cross_encoder_rerank(query, results['documents'][0], top_k=final_top_k)

    # Pipeline with decomposition
    aggregated = search_with_decomposition(query, top_k_per_subquery=5)

    # Rerank with the original query over the aggregated docs
    reranked = cross_encoder_rerank(
        query=query,
        documents=aggregated["documents"],
        top_k=final_top_k,
    )

    return reranked


# End-to-end usage
def answer_query(query: str) -> str:
    relevant_docs = smart_pipeline(query)
    answer = generate_answer(query, relevant_docs)
    return answer

QUERIES_PER_DAY = 5000
DAYS_PER_MONTH = 30

# Costs by category (decomposition LLM calls + extra retrievals)
PCT_DECOMPOSED = 0.18 + 0.15 + 0.08  # = 0.41 (41% gets decomposed)

# 1 LLM call for the decompose decision (gpt-4o-mini, ~200 tokens in/out)
COST_DECOMPOSE_DECISION = 0.000050

# 4 extra retrievals (vs 1 in the simple pipeline)
COST_EXTRA_RETRIEVALS = 4 * 0.000010  # extra OpenAI embedding calls

cost_per_decomposed_query = COST_DECOMPOSE_DECISION + COST_EXTRA_RETRIEVALS
# But the decomposition decision runs for ALL queries (it decides whether to decompose)
# Only on simple ones does it skip the 4 extra retrievals

monthly_cost = (
    QUERIES_PER_DAY * DAYS_PER_MONTH * COST_DECOMPOSE_DECISION  # decision for all
    + QUERIES_PER_DAY * DAYS_PER_MONTH * PCT_DECOMPOSED * COST_EXTRA_RETRIEVALS
)
print(f"Extra monthly cost: ${monthly_cost:.2f}")

Recap and next step

What you learned:

• Query decomposition splits complex queries (comparative, multi-step, causal) into independent sub-queries.
• The LLM decides whether to decompose (structured output with is_decomposable). A single call.
• Search sub-queries in parallel (ThreadPoolExecutor) to keep latency reasonable.
• After aggregating, rerank with the ORIGINAL query (not the sub-queries) to favor holistic docs.
• Cap the final context sent to the LLM at 5-10 docs after the rerank — more causes "lost in the middle".
• Main trap: decomposing simple queries. The structured output with automatic detection prevents it.
• Decomposition + rerank + dynamic skip is the combination that delivers quality without blowing up the cost.

Checkpoint: before moving on, you should be able to:

• Identify the three classes of query that benefit from decomposition.
• Implement decomposition with structured outputs and automatic detection.
• Design a cascading pipeline with decomposition + rerank and a dynamic skip.

Next capsule: 06 — HyDE (Hypothetical Document Embeddings).

Up to here, query optimization techniques manipulate the query as text. HyDE is different: instead of searching with the query's embedding, you generate a hypothetical answer with an LLM and embed that answer. What you look for in the corpus is "documents similar to the kind of answer an expert would give" — semantically closer to the real docs than the raw question is. It's the most sophisticated technique in the module, useful when nothing else is enough.

Resources

1. LangChain — Multi-Query and Decomposition — Implemented patterns
2. LlamaIndex — Sub Question Query Engine — Official decomposition implementation
3. Anthropic — Multi-Step Reasoning — Complementary patterns
4. OpenAI — Structured Outputs — For a robust implementation
5. Lost in the Middle Paper — Why limiting context matters
6. Self-Ask Prompting (Press et al., 2022) — Decomposition via chain-of-thought

Estimated time: 30-35 minutes Next: 06-hyde.md