RAG Pipeline Architecture for the Enterprise: Five Layers Beyond the Basic Chatbot

from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import Chroma
from langchain.retrievers import EnsembleRetriever
from langchain_cohere import CohereRerank
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.schema import Document
from typing import List
import os

# --- Document corpus (in production, load from your vector DB and document store) ---
def load_documents() -> List[Document]:
    """Load documents from your data source.
    In production: use LangChain document loaders for PDF, DOCX, SharePoint, etc.
    """
    # Placeholder — replace with your actual document loading pipeline
    return [
        Document(
            page_content="Q3 2025 gross margin declined 4.2pp due to the migration to Vendor B for PCB components, which carried a 12% price premium over legacy Vendor A.",
            metadata={"source": "board_report_q3_2025.pdf", "page": 7, "section": "Margin Analysis"}
        ),
        Document(
            page_content="The supplier transition from Vendor A to Vendor B was completed in July 2025 following a quality audit that found Vendor A's yield rate had dropped below the 98% contractual threshold.",
            metadata={"source": "supply_chain_summary_2025.pdf", "page": 3, "section": "Supplier Changes"}
        ),
        Document(
            page_content="Q3 2025 revenue grew 8% year-over-year to $142M. Operating expenses increased 11% primarily due to headcount additions in the engineering division.",
            metadata={"source": "board_report_q3_2025.pdf", "page": 4, "section": "Financial Highlights"}
        ),
    ]

def build_hybrid_rag_chain(documents: List[Document]) -> RetrievalQA:
    embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
    llm = ChatOpenAI(model="gpt-4o", temperature=0)

    # Dense vector retriever
    vectorstore = Chroma.from_documents(documents, embeddings)
    dense_retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 10}  # Retrieve 10 candidates before reranking
    )

    # Sparse BM25 retriever
    bm25_retriever = BM25Retriever.from_documents(documents)
    bm25_retriever.k = 10

    # Ensemble: combine dense + sparse with equal weighting
    # alpha=0.5 gives equal weight to both; tune toward 0.7 for dense-heavy corpora
    ensemble_retriever = EnsembleRetriever(
        retrievers=[bm25_retriever, dense_retriever],
        weights=[0.4, 0.6]  # Slight preference for semantic over keyword
    )

    # Cohere reranker: re-score top-10 candidates, return top-4
    cohere_reranker = CohereRerank(
        model="rerank-english-v3.0",
        top_n=4,
        cohere_api_key=os.environ["COHERE_API_KEY"]
    )

    # Wrap ensemble retriever with contextual compression (reranking)
    compression_retriever = ContextualCompressionRetriever(
        base_compressor=cohere_reranker,
        base_retriever=ensemble_retriever
    )

    # Build QA chain with source citation enforcement
    system_prompt = (
        "You are an enterprise document analyst. Answer questions using ONLY the provided context. "
        "Cite the source document name for every factual claim using [Source: filename] notation. "
        "If the context does not contain sufficient information, respond with: "
        "'The available documentation does not provide enough information to answer this question reliably.'"
    )

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=compression_retriever,
        return_source_documents=True,
        chain_type_kwargs={
            "prompt": None  # Use default; in production, inject system_prompt via PromptTemplate
        }
    )

    return chain

def main():
    documents = load_documents()
    chain = build_hybrid_rag_chain(documents)

    query = "What was the impact of the 2025 supplier change on Q3 margins?"
    result = chain.invoke({"query": query})

    print(f"Query: {query}")
    print(f"\nAnswer:\n{result['result']}")
    print("\nSource Documents:")
    for doc in result["source_documents"]:
        print(f"  - {doc.metadata['source']} (page {doc.metadata.get('page', 'N/A')})")

if __name__ == "__main__":
    main()