Multi-Stage RAG with Azure AI Search: Real Production Lessons & Architecture

1. The Naive RAG Trap in Enterprise Systems

Almost every proof-of-concept RAG pipeline begins the same way: load raw PDFs, split them using an arbitrary chunk size (e.g., 1,000 characters with 200 character overlap), embed chunks into a vector database with cosine similarity, and dump the top-5 results into the prompt context.

In high-stakes production domains—such as legal compliance, ERP data mining, and technical medical records—this naive approach fails predictably:

• Loss of Document Hierarchy: Section headings, table schemas, footnotes, and regulatory clauses get severed from their parent context.
• Keyword & Exact-Match Blindness: Dense vector embeddings struggle with acronyms, part numbers, section codes (e.g., Section 409A(b)(2)), and exact currency figures.
• Lost in the Middle: Long context windows suffer significant recall degradation when the critical fact is buried in the middle third of the retrieved context.
• Hallucinatory Extrapolation: When retrieved chunks are only tangentially related, LLMs fabricate connective reasoning to satisfy the user prompt.

2. The 5-Layer Multi-Stage Retrieval Architecture

To solve this, I engineered a production-grade 5-layer retrieval system deployed on Azure AI Search and orchestrated through Python FastAPI.

Stage 1: Query Reformulation & Hypothetical Document Embeddings (HyDE)

Raw user queries are frequently ambiguous, conversational, or under-specified. Before executing a search:

1. Multi-Query Expansion: The system generates 3 orthogonal search representations (e.g., keyword-dense, conceptual, and hypothetical answer text).
2. Entity Extraction: Structured filters (e.g., date ranges, document IDs, jurisdiction) are parsed into deterministic Azure OData filter clauses (e.g., $filter=doc_type eq 'contract' and fiscal_year ge 2023).

Stage 2: Dual Hybrid Retrieval (Dense Vector + BM25 Okapi)

I execute parallel retrieval combining dense vector semantic search using text-embedding-3-large (3,072 dimensions) with lexical keyword matching using BM25 Okapi tokenized across custom analyzer dictionaries.

# Azure AI Search Hybrid Query with Vector and Lexical Search
from azure.search.documents import SearchClient
from azure.search.documents.models import VectorizedQuery

async def execute_hybrid_search(search_client: SearchClient, query_text: str, query_vector: list[float], filter_expr: str):
    vector_query = VectorizedQuery(
        vector=query_vector,
        k_nearest_neighbors=50,
        fields="embedding_vector"
    )
    
    results = await search_client.search(
        search_text=query_text,
        vector_queries=[vector_query],
        filter=filter_expr,
        query_type="semantic",
        semantic_configuration_name="legal-doc-semantic-config",
        top=40
    )
    return [doc async for doc in results]

Stage 3: Reciprocal Rank Fusion (RRF)

To combine results from disparate search mechanisms with different score distributions, the pipeline calculates the Reciprocal Rank Fusion score:

$$RRF(d) = \sum_{m \in M} \frac{1}{k + r_m(d)}$$

Where $r_m(d)$ is the rank of document $d$ in retrieval method $m$, and $k$ is a smoothing constant set to $60$. This yields a robust candidate set of the top 40 passages.

Stage 4: Cross-Encoder Deep Reranking

Bi-encoders (embedding vectors) compute document representations independently. Cross-encoders pass both the query and candidate chunk jointly through all transformer layers, capturing nuanced token-level cross-attention.

Applying Cohere Rerank v3 over the top 40 candidates prunes the set down to the top 5–6 highest-signal chunks, discarding false positives that scored high on vector similarity but lacked specific semantic relevance.

Stage 5: Grounded Generation & Citation Verification

The synthesized prompt enforces strict citation attribution. Every generated sentence must map to a discrete source chunk identifier:

{
  "answer": "Under Section 14.2, early termination requires 30 days prior written notice accompanied by payment of unamortized installation fees.",
  "citations": [
    {
      "source_document": "Master_Services_Agreement_2024.pdf",
      "page_number": 18,
      "chunk_id": "doc_8492_p18_c3",
      "verbatim_snippet": "either party may terminate upon thirty (30) days prior written notice subject to reimbursement of unamortized setup expenses"
    }
  ],
  "confidence_score": 0.97
}

3. Document Ingestion & Hierarchical Chunking

Standard fixed-character chunking corrupts tables and multi-level lists. I implemented Markdown-aware semantic chunking:

• PDFs are parsed via Azure Document Intelligence to preserve tabular structures as native HTML tables.
• Document text is segmented by heading levels (H1 > H2 > H3).
• Small child chunks (200 tokens) are indexed for fine-grained retrieval, while each chunk maintains a pointer to its surrounding 1,200-token parent context section, passed to the LLM during generation.

4. Automated Evaluation & Measured Results

Using the Ragas automated evaluation framework across a curated golden evaluation dataset of 250 enterprise queries, the multi-stage pipeline outperformed naive RAG across every reliability metric:

5. Key Takeaways for Production Implementations

1. Hybrid is non-negotiable: Dense vectors alone miss domain vocabulary; lexical search alone misses conceptual synonyms.
2. Reranking delivers the highest ROI: Adding a cross-encoder reranker to your pipeline provides a bigger accuracy leap than upgrading your base LLM from 8B to 70B parameters.
3. Deterministic metadata filters save compute: Use LLM query extractors to pre-filter by organization, date, and document type before executing vector queries.