Complete RAG Pipeline

This is the complete end-to-end pipeline. Drop your PDFs into ./documents/, set your API key, run it.

What this gives you:
✓ PDF loading with page metadata
✓ Recursive text splitting with overlap
✓ BGE-large embeddings (best open source)
✓ ChromaDB with persistence (no re-indexing on restart)
✓ Hybrid BM25 + Dense retrieval (catches both keywords and semantics)
✓ Cross-encoder reranking (accuracy boost)
✓ Gemini LLM with grounded prompt
✓ Source citations in every answer

Retrieve 10 → Rerank to 3 → Generate from 3.
That's the formula.

1"""
2deepdocs — Complete LangChain RAG Pipeline
3──────────────────────────────────────────
4pip install langchain langchain-community langchain-huggingface
5pip install langchain-google-genai chromadb sentence-transformers
6pip install pypdf rank-bm25
7"""
8
9from pathlib import Path
10from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
11from langchain.text_splitter import RecursiveCharacterTextSplitter
12from langchain_huggingface import HuggingFaceEmbeddings
13from langchain_community.vectorstores import Chroma
14from langchain_community.retrievers import BM25Retriever
15from langchain.retrievers import EnsembleRetriever
16from langchain_core.prompts import ChatPromptTemplate
17from langchain_core.output_parsers import StrOutputParser
18from langchain_core.runnables import RunnableLambda, RunnablePassthrough
19from langchain_google_genai import ChatGoogleGenerativeAI
20from sentence_transformers import CrossEncoder
21
22# ── Config ────────────────────────────────────────────────────────
23GOOGLE_API_KEY = "YOUR_GOOGLE_API_KEY"
24DOCS_PATH      = "./documents/"
25CHROMA_PATH    = "./chroma_db"
26EMBED_MODEL    = "BAAI/bge-large-en-v1.5"
27RERANK_MODEL   = "cross-encoder/ms-marco-MiniLM-L-6-v2"
28CHUNK_SIZE     = 800
29CHUNK_OVERLAP  = 150
30RETRIEVE_K     = 10
31RERANK_TOP_N   = 3
32
33
34def load_documents(path: str):
35    print("📂 Loading documents...")
36    loader = DirectoryLoader(path, glob="**/*.pdf",
37                             loader_cls=PyPDFLoader, show_progress=True)
38    docs = loader.load()
39    print(f"   ✓ {len(docs)} pages loaded")
40    return docs
41
42
43def split_documents(docs):
44    print("✂️  Splitting into chunks...")
45    splitter = RecursiveCharacterTextSplitter(
46        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP,
47        separators=["\n\n", "\n", ". ", " ", ""],
48    )
49    chunks = splitter.split_documents(docs)
50    for i, c in enumerate(chunks):
51        c.metadata["chunk_id"] = i
52    print(f"   ✓ {len(chunks)} chunks created")
53    return chunks
54
55
56def build_vectorstore(chunks, force_rebuild=False):
57    embeddings = HuggingFaceEmbeddings(
58        model_name=EMBED_MODEL,
59        encode_kwargs={"normalize_embeddings": True},
60    )
61    if Path(CHROMA_PATH).exists() and not force_rebuild:
62        print("📦 Loading existing vector store...")
63        vs = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
64    else:
65        print("🔢 Embedding and indexing chunks...")
66        vs = Chroma.from_documents(chunks, embeddings,
67                                   persist_directory=CHROMA_PATH)
68    print(f"   ✓ {vs._collection.count()} chunks in store")
69    return vs, embeddings
70
71
72def build_retriever(vs, chunks):
73    bm25  = BM25Retriever.from_documents(chunks)
74    bm25.k = RETRIEVE_K
75    dense = vs.as_retriever(search_kwargs={"k": RETRIEVE_K})
76    return EnsembleRetriever(retrievers=[bm25, dense], weights=[0.4, 0.6])
77
78
79def build_reranker():
80    print("⚡ Loading cross-encoder...")
81    model = CrossEncoder(RERANK_MODEL, max_length=512)
82    print("   ✓ Cross-encoder ready")
83    return model
84
85
86def rerank(query, docs, model, top_n=RERANK_TOP_N):
87    pairs  = [(query, doc.page_content) for doc in docs]
88    scores = model.predict(pairs)
89    ranked = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
90    return [doc for _, doc in ranked[:top_n]]
91
92
93def build_chain(retriever, reranker):
94    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash",
95                                 google_api_key=GOOGLE_API_KEY, temperature=0.1)
96    prompt = ChatPromptTemplate.from_template("""
97Answer using ONLY the context below. If insufficient, say so.
98
99Context:
100{context}
101
102Question: {question}
103Answer:""")
104
105    def retrieve_rerank_format(inputs):
106        q    = inputs["question"]
107        docs = retriever.invoke(q)
108        top  = rerank(q, docs, reranker)
109        return "\n\n---\n\n".join([
110            f"[{doc.metadata.get('source','?')}, p{doc.metadata.get('page','?')}]"
111            f"\n{doc.page_content}" for doc in top
112        ])
113
114    return (
115        {"context": RunnableLambda(retrieve_rerank_format),
116         "question": RunnablePassthrough() | (lambda x: x["question"])}
117        | prompt | llm | StrOutputParser()
118    )
119
120
121if __name__ == "__main__":
122    docs      = load_documents(DOCS_PATH)
123    chunks    = split_documents(docs)
124    vs, _     = build_vectorstore(chunks)
125    retriever = build_retriever(vs, chunks)
126    reranker  = build_reranker()
127    chain     = build_chain(retriever, reranker)
128
129    print("\n" + "═" * 60)
130    print("deepdocs RAG ready. Ask anything.")
131    print("═" * 60 + "\n")
132
133    while True:
134        q = input("Question (q to quit): ").strip()
135        if q.lower() == "q":
136            break
137        print("\n" + chain.invoke({"question": q}) + "\n")