Boosting AI Agent Reliability with Parallel Query Expansion

Parallel Query Expansion Design

To increase retrieval recall, the system first asks an LLM to generate a set of diversified search queries before any document lookup. The generated queries include a hypothetical document (HyDE), a short list of sub‑questions, and a list of core keywords.

Pydantic Model for Expanded Queries

from langchain_core.pydantic_v1 import BaseModel, Field
from typing import List

class ExpandedQueries(BaseModel):
    """Defines a set of expanded queries to improve retrieval recall"""
    # Paragraph‑length hypothetical document that directly answers the user's question
    hypothetical_document: str = Field(
        description="A generated, paragraph‑length hypothetical document that directly answers the user's question, which will be used for semantic search.",
        alias="hyde_query"
    )
    # 2‑3 smaller, more specific questions derived from the original query
    sub_questions: List[str] = Field(
        description="A list of 2‑3 smaller, more specific questions that break down the original query."
    )
    # 3‑5 core keywords or entities for lexical search
    keywords: List[str] = Field(
        description="A list of 3‑5 core keywords and entities extracted from the user's query."
    )

Graph Nodes

Query Expansion Node – receives the original user question, invokes a structured‑output chain, and fills the ExpandedQueries model.

query_expansion_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a query expansion specialist. Your goal is to transform a user's question into a diverse set of search queries to maximize retrieval recall. Generate a hypothetical document, sub‑questions, and keywords."),
    ("human", "Please expand the following question: {question}")
])

query_expansion_chain = query_expansion_prompt | llm.with_structured_output(ExpandedQueries)

def query_expansion_node(state: RAGGraphState):
    print("--- [Expander] Generating parallel queries... ---")
    expanded_queries = query_expansion_chain.invoke({"question": state['original_question']})
    return {"expanded_queries": expanded_queries}

Retrieval Node – aggregates all generated queries, dispatches them concurrently with ThreadPoolExecutor (max_workers=5), collects returned documents, and deduplicates by content.

from concurrent.futures import ThreadPoolExecutor

def retrieval_node(state: RAGGraphState):
    print("--- [Retriever] Executing parallel searches... ---")
    all_queries = []
    expanded = state['expanded_queries']
    all_queries.append(expanded.hypothetical_document)
    all_queries.extend(expanded.sub_questions)
    all_queries.extend(expanded.keywords)

    all_docs = []
    with ThreadPoolExecutor(max_workers=5) as executor:
        results = executor.map(retriever.invoke, all_queries)
    for docs in results:
        all_docs.extend(docs)
    unique_docs = {doc.page_content: doc for doc in all_docs}.values()
    print(f"--- [Retriever] Found {len(unique_docs)} unique documents from {len(all_queries)} queries. ---")
    return {"retrieved_docs": list(unique_docs)}

Workflow Assembly

from langgraph.graph import StateGraph, END

workflow = StateGraph(RAGGraphState)
workflow.add_node("expand_queries", query_expansion_node)
workflow.add_node("retrieve_docs", retrieval_node)
workflow.add_node("generate_answer", generation_node)
workflow.set_entry_point("expand_queries")
workflow.add_edge("expand_queries", "retrieve_docs")
workflow.add_edge("retrieve_docs", "generate_answer")
workflow.add_edge("generate_answer", END)

Experimental Comparison

Both a simple RAG pipeline (single user query) and the advanced pipeline (with parallel query expansion) are run on the ambiguous query:

user_query = "How do modern AI systems get so big and fast at the same time? I've heard about attention but I'm not sure how it's optimized."

Simple RAG retrieves one document about multi‑headed attention. Advanced RAG retrieves three documents covering FlashAttention, Mixture of Experts, and multi‑headed attention.

--- Simple RAG Retrieved 1 document(s) ---
1. **Multi‑headed Attention Mechanism**: The core component of the Transformer architecture ...

--- Advanced RAG Retrieved 3 document(s) ---
1. **FlashAttention Optimization**: ... an I/O‑aware algorithm that reduces read/write operations ...
2. **Mixture of Experts (MoE) Layers**: ... a router network dynamically selects a small subset of 'expert' sub‑networks ...
3. **Multi‑headed Attention Mechanism**: The core component of the Transformer architecture ...

Analysis

The query_expansion_node introduces missing technical terms (e.g., "Mixture of Experts", "FlashAttention") that bridge the semantic gap between the user’s vague wording and the terminology used in the knowledge base.

By executing all expanded queries in parallel, the advanced system captures additional relevant documents, increasing recall from 1 to 3 documents for the same query.

Higher recall provides richer context for the answer generation step, resulting in a more complete and technically accurate final answer.

All code, models, and experiments are available in the GitHub repository: https://github.com/FareedKhan-dev/agentic-parallelism