Parallel Query Expansion: Boosting Reliability in Agentic AI Systems

Parallel Query Expansion for Retrieval Augmented Generation

Vocabulary mismatch is a common failure mode in autonomous RAG pipelines. A user query such as “How do modern AI systems get so big and fast?” may miss technical terms like “Mixture of Experts” or “FlashAttention”. Parallel Query Expansion mitigates this by having an LLM generate a set of diversified queries before retrieval.

Expanded query schema

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

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

The LLM is prompted with a system message to act as a query‑expansion specialist and returns an ExpandedQueries instance in a single structured call.

def query_expansion_node(state: RAGGraphState):
    """Generate a set of expanded queries from the original question."""
    print("--- [Expander] Generating parallel queries... ---")
    expanded_queries = query_expansion_chain.invoke({"question": state['original_question']})
    return {"expanded_queries": expanded_queries}

Parallel retrieval

The retrieval node collects all generated queries (hypothetical document, sub‑questions, keywords) and executes them concurrently with ThreadPoolExecutor, then deduplicates the results.

def retrieval_node(state: RAGGraphState):
    """Execute all expanded queries in parallel."""
    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 typing import TypedDict, List, Optional
from langchain_core.documents import Document

class RAGGraphState(TypedDict):
    original_question: str
    expanded_queries: Optional[ExpandedQueries]
    retrieved_docs: List[Document]
    final_answer: str

# Prompt definition
query_expansion_prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a query expansion specialist. 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)

# Graph construction
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

A side‑by‑side experiment runs a simple RAG pipeline (single query) and the advanced parallel‑query RAG on the same vague user query. The simple pipeline retrieves one document about multi‑headed attention. The advanced pipeline retrieves three documents covering FlashAttention, Mixture of Experts, and multi‑headed attention.

# Simple RAG
simple_retrieved_docs = retriever.invoke(user_query)

# Advanced RAG (graph execution)
advanced_rag_result = workflow.invoke({"original_question": user_query})

Output shows:

--- Simple RAG Retrieved 1 document ---
**Multi-headed Attention Mechanism** ...

--- Advanced RAG Retrieved 3 documents ---
**FlashAttention Optimization** ...
**Mixture of Experts (MoE) Layers** ...
**Multi-headed Attention Mechanism** ...

Analysis indicates higher recall: the parallel‑query system captures missing technical terms, provides richer context, and enables more accurate answers.

All code and notebooks are hosted in the GitHub repository:

https://github.com/FareedKhan-dev/agentic-parallelism

References:

[1] Building the 14 Key Pillars of Agentic AI – https://levelup.gitconnected.com/building-the-14-key-pillars-of-agentic-ai-229e50f65986

[2] Speculative execution – https://en.wikipedia.org/wiki/Speculative_execution

[3] Redundant execution – https://developer.arm.com/community/arm-community-blogs/b/embedded-and-microcontrollers-blog/posts/comparing-lock-step-redundant-execution-versus-split-lock-technologies

[4] Agentic Parallelism repository – https://github.com/FareedKhan-dev/agentic-parallelism