07. Query Expansion

Query expansion increases recall by generating multiple queries that explore different aspects of a question, then retrieving documents for each and combining results. This chapter covers techniques ranging from simple augmentation to sophisticated multi-query strategies.

Multi-Query Retrieval

The straightforward approach: generate multiple queries from the original, retrieve documents for each, and combine results:

from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough

def multi_query_retrieval(query, retrievers, llm, combiner):
    """
    Generate multiple queries, retrieve for each, combine results.
    
    Args:
        query: Original user query
        retrievers: Dict of retrievers by query type
        llm: LLM for query generation
        combiner: Function to merge and rank results
    """
    # Generate diverse queries
    query_gen_prompt = f"""Given this user query: {query}
    
    Generate 5 different search queries that:
    1. Addresses the core topic from the original question
    2. Explores a specific sub-topic or angle
    3. Uses different terminology
    4. Asks from a different stakeholder perspective
    5. Focuses on a related but distinct concept
    
    Output each query on a new line."""
    
    generated_queries = llm.invoke(query_gen_prompt)
    queries = parse_query_list(generated_queries.content)
    
    # Retrieve for each query
    query_results = {}
    for i, q in enumerate(queries):
        # Apply query-specific retrieval strategy
        strategy = query_strategy(q)  # simplistic/custom logic
        if strategy == 'keyword':
            results = retrievers['keyword'].get_relevant_documents(q)
        elif strategy == 'semantic':
            results = retrievers['semantic'].get_relevant_documents(q)
        else:
            results = retrievers['default'].get_relevant_documents(q)
        
        query_results[f"query_{i}"] = {
            'query': q,
            'results': results
        }
    
    # Combine and deduplicate
    combined = combiner(query_results)
    
    return {
        'queries': queries,
        'combined_results': combined
    }

Sub-Query Decomposition

Complex questions contain multiple implicit sub-questions. Decomposition breaks them into explicit queries:

def decompose_query(query, llm):
    """Decompose compound queries into atomic sub-queries."""
    decomposition_prompt = f"""Analyze this query and break it into atomic sub-questions.
Each sub-question should be answerable by retrieving a single document or passage.
If the original query can be fully answered by one sub-question, return just that.

Original query: {query}

Sub-questions:"""

    response = llm.invoke(decomposition_prompt)
    sub_questions = parse_query_list(response.content)
    return sub_questions

# Example
original = "What is the vacation policy and how does it compare to sick leave?"
decomposed = decompose_query(original)
# Output:
# - "What is the employee vacation policy?"
# - "What is the sick leave policy?"
# - "How do sick leave benefits compare to vacation benefits?"

After decomposition, retrieve for each sub-question and combine. This ensures that each aspect of the compound query gets dedicated retrieval effort.

Step-Back Query Generation

Sometimes narrowing retrieval is the problem, not broadening it. A "step-back" query generalizes to the broader topic:

def generate_step_back(query, llm):
    """Generate a more general query that captures the broader topic."""
    step_back_prompt = f"""Given this query: {query}

What is the broader topic or principle this question is about?
Generate a more general query that would retrieve background information.
The step-back query should help understand the domain context.

Step-back query:"""

    response = llm.invoke(step_back_prompt)
    return response.content.strip()

# Example
original = "Why is our Q3 revenue lower than expected?"
step_back = generate_step_back(original)
# Output: "Q3 revenue reporting and analysis methodology"

Step-back queries retrieve foundational context that improves answer ability. They're particularly useful for analytical questions where understanding requires domain knowledge first.

Selective Expansion

Not every query benefits from expansion. Use expansion selectively:

Expand when: Query is compound (multiple questions), query is ambiguous (could mean several things), query topic is broad (many relevant documents), evaluation shows low recall.

Don't expand when: Query is already precise, expansion significantly increases latency, unique terminology already matches documents, query is time-sensitive (expansion adds latency).

def selective_expansion(query, classifier, llm):
    """Conditionally expand based on query characteristics."""
    # Classify query type
    query_type = classifier.predict([query])[0]
    
    if query_type == 'simple':
        # Direct retrieval, no expansion
        return {'queries': [query], 'strategy': 'direct'}
    
    elif query_type == 'compound':
        # Decompose into sub-queries
        sub_queries = decompose_query(query, llm)
        return {'queries': sub_queries, 'strategy': 'decompose'}
    
    elif query_type == 'ambiguous':
        # Multiple alternative expansions
        expansions = expand_query(query, llm)
        return {'queries': expansions, 'strategy': 'expand'}
    
    else:  # broad
        # Step-back plus specific retrieval
        step_back_q = generate_step_back(query, llm)
        return {'queries': [step_back_q, query], 'strategy': 'step_back'}

Train a simple classifier on query patterns to route to appropriate strategies:

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

# Example training data
queries = [
    ("What is the CEO's name?", "simple"),
    ("Tell me about the project status", "simple"),
    ("What are the vacation and sick leave policies?", "compound"),
    ("What is the approved vendor list?", "simple"),
    ("How do I dispute a charge?", "simple"),
    ("Compare benefits across all departments", "broad"),
]

train_queries, train_labels = zip(*queries)
vectorizer = TfidfVectorizer(ngram_range=(1, 2))
X_train = vectorizer.fit_transform(train_queries)

classifier = LogisticRegression()
classifier.fit(X_train, train_labels)