Gemini 2.0 with LangChain: Production RAG Pipeline 2026

# Use uv for fast installs (pip works too)
uv add langchain langchain-google-genai langchain-community chromadb tiktoken

# For LangSmith tracing (recommended for production)
uv add langsmith

python -c "import langchain_google_genai; print(langchain_google_genai.__version__)"
# Expected: 2.x.x

# .env
GOOGLE_API_KEY=your_key_here

# Optional but strongly recommended
LANGCHAIN_TRACING_V2=true
LANGCHAIN_API_KEY=your_langsmith_key
LANGCHAIN_PROJECT=rag-production

import os
from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings

load_dotenv()

# Gemini 2.0 Flash: fast, cheap, 1M context
llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    temperature=0.1,          # Low temp for factual RAG answers
    max_output_tokens=2048,
    # Convert_system_message_to_human: required for older Gemini versions
    # Gemini 2.0 supports system messages natively — leave this out
)

# text-embedding-004 outperforms text-embedding-003 on MTEB benchmarks
embeddings = GoogleGenerativeAIEmbeddings(
    model="models/text-embedding-004",
    task_type="retrieval_document",  # Optimizes embeddings for retrieval, not similarity
)

from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, TextFileLoader

# Load documents from a local directory
loader = DirectoryLoader(
    "./docs",
    glob="**/*.txt",
    loader_cls=TextFileLoader,
    show_progress=True,
)
raw_docs = loader.load()

# Chunk size 1000 with 200 overlap is a solid default for most content
# Increase chunk_size for technical docs; decrease for Q&A datasets
splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    length_function=len,
    add_start_index=True,   # Adds char offset to metadata — useful for debugging
)
docs = splitter.split_documents(raw_docs)

print(f"Split {len(raw_docs)} documents into {len(docs)} chunks")

# Persist to disk so you don't re-embed on every restart
vectorstore = Chroma.from_documents(
    documents=docs,
    embedding=embeddings,
    persist_directory="./chroma_db",
    collection_name="production_rag",
)

from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

# Load an existing vectorstore (skip re-embedding)
vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings,
    collection_name="production_rag",
)

# MMR retrieval reduces redundancy vs plain similarity search
# fetch_k=20 retrieves 20 candidates, then reranks to k=5
retriever = vectorstore.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 5,
        "fetch_k": 20,
        "lambda_mult": 0.7,  # 1.0 = pure similarity, 0.0 = pure diversity
    },
)

# System prompt: be explicit about what to do when context is insufficient
system_prompt = """You are a helpful assistant answering questions based only on the provided context.

Context:
{context}

Rules:
- Answer only from the context above
- If the context doesn't contain the answer, say "I don't have enough information to answer that"
- Cite the source document when possible
- Be concise: 2–4 sentences unless asked for more
"""

prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])

# create_stuff_documents_chain concatenates retrieved docs into {context}
question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)

import asyncio
from tenacity import retry, stop_after_attempt, wait_exponential
from langchain_core.runnables import RunnableConfig

@retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10),
    # Retry on Google API rate limit errors (429) and transient 5xx
    reraise=True,
)
async def stream_rag_response(query: str):
    """Stream tokens from the RAG chain with automatic retry."""
    config = RunnableConfig(
        tags=["production", "rag"],       # Appears in LangSmith trace
        metadata={"query_length": len(query)},
    )

    full_answer = ""

    async for chunk in rag_chain.astream({"input": query}, config=config):
        # Retrieval chain streams: first the docs, then the answer tokens
        if "answer" in chunk:
            token = chunk["answer"]
            full_answer += token
            print(token, end="", flush=True)

    print()  # Newline after stream completes
    return full_answer


# Run it
asyncio.run(stream_rag_response("What are the main features of the product?"))

from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever

# Query-optimized embeddings for the retriever
query_embeddings = GoogleGenerativeAIEmbeddings(
    model="models/text-embedding-004",
    task_type="retrieval_query",  # Different from document embeddings above
)

# Hybrid retrieval: BM25 (keyword) + vector (semantic) = better coverage
bm25_retriever = BM25Retriever.from_documents(docs)
bm25_retriever.k = 5

vector_retriever = vectorstore.as_retriever(
    search_type="mmr",
    search_kwargs={"k": 5, "fetch_k": 20},
)

# EnsembleRetriever merges results with Reciprocal Rank Fusion
hybrid_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.4, 0.6],  # Favor semantic search; adjust based on your content
)

# Rebuild the chain with hybrid retrieval
rag_chain = create_retrieval_chain(
    hybrid_retriever,
    question_answer_chain,
)

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import json

app = FastAPI()

class QueryRequest(BaseModel):
    question: str
    max_length: int = 500


@app.post("/query")
async def query_rag(request: QueryRequest):
    if not request.question.strip():
        raise HTTPException(status_code=400, detail="Question cannot be empty")

    async def generate():
        async for chunk in rag_chain.astream({"input": request.question}):
            if "answer" in chunk:
                # Server-Sent Events format for easy frontend consumption
                yield f"data: {json.dumps({'token': chunk['answer']})}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")


@app.get("/health")
async def health():
    return {"status": "ok", "model": "gemini-2.0-flash"}

uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4

import requests

response = requests.post(
    "http://localhost:8000/query",
    json={"question": "What is this product used for?"},
    stream=True,
)

for line in response.iter_lines():
    if line and line.startswith(b"data: ") and line != b"data: [DONE]":
        data = json.loads(line[6:])
        print(data["token"], end="", flush=True)

open https://smith.langchain.com
# Navigate to your project: rag-production
# Inspect retrieval latency, chunk quality, and token counts per query