Reranking retrieved documents in a LangChain RAG pipeline

Reranking is the single highest-leverage improvement you can make to a LangChain RAG pipeline after basic retrieval works. The retriever fetches candidates — often 20 to 50 chunks — but the LLM only has context window for a fraction. A cross-encoder reranker scores each candidate against the query with full attention, pushing the truly relevant chunks to the top. This guide walks through adding reranking to a LangChain RAG pipeline using both open-source cross-encoders and the Cohere rerank API, with verification steps at each stage.

Step 1: Set up the baseline retrieval pipeline

Start with a working vector store and retriever. If you already have this, skip to step 2. We’ll use Chroma with OpenAI embeddings for the example, but the reranking layer is vector-store agnostic.

# baseline_rag.py
import os
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader

# Load and split documents
loader = TextLoader("data/your_docs.txt")
docs = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = splitter.split_documents(docs)

# Embed and store
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(splits, embeddings, persist_directory="./chroma_db")
retriever = vectorstore.as_retriever(search_kwargs={"k": 20})  # fetch more for reranking

# Basic RAG chain
prompt = ChatPromptTemplate.from_template("""Answer the question using only the context below.

Context:
{context}

Question: {question}
""")

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

def format_docs(docs):
    return "\n\n".join(d.page_content for d in docs)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# Test
print(rag_chain.invoke("What is the refund policy for enterprise customers?"))

Verify it works: run the script and confirm you get a coherent answer. The retriever returns 20 chunks; the LLM sees all of them. This is your baseline.

Step 2: Add a cross-encoder reranker with LangChain’s built-in compressor

LangChain ships ContextualCompressionRetriever and CrossEncoderReranker in langchain.retrievers.document_compressors. The cross-encoder runs locally via sentence-transformers — no API key needed.

# rerank_local.py
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder

# Use a strong open-source cross-encoder
model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
compressor = CrossEncoderReranker(model=model, top_n=5)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=retriever  # from step 1
)

# Swap into the chain
rag_chain_reranked = (
    {"context": compression_retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

print(rag_chain_reranked.invoke("What is the refund policy for enterprise customers?"))

What changed: base_retriever still fetches 20 chunks. The compressor passes each (query, chunk) pair through the cross-encoder, scores them, and returns only the top 5. The LLM now sees fewer, higher-signal chunks.

Verify: compare the retrieved chunk IDs before and after. Add a quick debug wrapper:

def debug_retriever(retriever, query):
    docs = retriever.invoke(query)
    for i, d in enumerate(docs):
        print(f"[{i}] {d.metadata.get('source', 'unknown')} :: {d.page_content[:120]}...")
    return docs

print("=== Baseline ===")
debug_retriever(retriever, "refund policy enterprise")

print("\n=== Reranked ===")
debug_retriever(compression_retriever, "refund policy enterprise")

You should see different ordering and fewer results from the reranked version.

Step 3: Use Cohere rerank for production latency and quality

Local cross-encoders work well for prototypes but add GPU memory and latency. Cohere’s rerank endpoint is faster at scale and often higher quality. LangChain has a first-party integration.

pip install langchain-cohere

# rerank_cohere.py
from langchain_cohere import CohereRerank
from langchain.retrievers import ContextualCompressionRetriever

# Requires COHERE_API_KEY in env
cohere_reranker = CohereRerank(
    model="rerank-v3.5",
    top_n=5,
    # Optional: pass user_id for analytics, or max_tokens_per_doc
)

cohere_compression_retriever = ContextualCompressionRetriever(
    base_compressor=cohere_reranker,
    base_retriever=retriever
)

rag_chain_cohere = (
    {"context": cohere_compression_retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

print(rag_chain_cohere.invoke("What is the refund policy for enterprise customers?"))

Verify: same debug function. Cohere returns a relevance_score in each document’s metadata — inspect it:

docs = cohere_compression_retriever.invoke("refund policy enterprise")
for d in docs:
    print(f"score={d.metadata.get('relevance_score'):.4f} :: {d.page_content[:100]}")

Scores near 1.0 are highly relevant; near 0.0 are noise. The compressor drops anything below the implicit threshold after sorting by score.

Step 4: Tune k and top_n for your corpus

Two knobs control the recall-precision tradeoff:

• search_kwargs={"k": N} on the base retriever — how many candidates to fetch from the vector store
• top_n=M on the reranker — how many to keep after scoring

Rule of thumb: set k to 20–50 for dense retrieval, top_n to 3–8 for the LLM context. Larger k improves recall but increases reranker latency linearly. Cohere bills per document reranked; local models burn GPU time.

Run a quick grid search on a held-out eval set:

# tune.py
import time
from statistics import mean

def evaluate(k, top_n, queries, expected_snippets):
    retriever = vectorstore.as_retriever(search_kwargs={"k": k})
    compressor = CohereRerank(model="rerank-v3.5", top_n=top_n)
    c_retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=retriever)
    
    latencies = []
    hits = 0
    for q, expected in zip(queries, expected_snippets):
        start = time.perf_counter()
        docs = c_retriever.invoke(q)
        latencies.append(time.perf_counter() - start)
        # Check if any returned chunk contains expected snippet
        if any(expected in d.page_content for d in docs):
            hits += 1
    return hits / len(queries), mean(latencies)

queries = [
    "refund policy enterprise",
    "SLA uptime guarantee",
    "data processing agreement",
]
expected = [
    "enterprise customers receive full refund within 30 days",
    "99.9% uptime",
    "DPA available on request",
]

for k in [10, 20, 30, 50]:
    for top_n in [3, 5, 8]:
        recall, latency = evaluate(k, top_n, queries, expected)
        print(f"k={k}, top_n={top_n} -> recall={recall:.2f}, latency={latency:.2f}s")

Pick the smallest k/top_n that hits your recall target. Typical sweet spot: k=30, top_n=5.

Step 5: Add rerank scoring to your observability pipeline

You need to know when reranking degrades — model drift, API changes, corpus shifts. Log the reranker scores alongside retrieval metrics.

# observable_rag.py
import json
import time
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

class RerankLogger(BaseCallbackHandler):
    def on_retriever_end(self, documents, *, run_id, parent_run_id, **kwargs):
        for i, doc in enumerate(documents):
            score = doc.metadata.get("relevance_score")
            if score is not None:
                print(json.dumps({
                    "event": "rerank_score",
                    "run_id": str(run_id),
                    "rank": i,
                    "score": score,
                    "source": doc.metadata.get("source"),
                    "chunk_id": doc.metadata.get("chunk_id"),
                }))

# Attach to the chain
from langchain_core.runnables import RunnableConfig

config = RunnableConfig(callbacks=[RerankLogger()])
result = rag_chain_cohere.invoke("refund policy enterprise", config=config)

Sample log line:

{"event": "rerank_score", "run_id": "abc-123", "rank": 0, "score": 0.987, "source": "policy_v4.md", "chunk_id": "chunk_12"}

Alert if median top-1 score drops below 0.7 for a rolling window — that signals corpus drift or query distribution shift.

Step 6: Handle edge cases — empty results, ties, and long documents

Rerankers can return empty lists if all scores are near zero. Guard against this:

from langchain_core.documents import Document

class SafeCompressionRetriever(ContextualCompressionRetriever):
    def _get_relevant_documents(self, query, *, run_manager):
        compressed = super()._get_relevant_documents(query, run_manager=run_manager)
        if not compressed:
            # Fall back to base retriever top-k without reranking
            return self.base_retriever.invoke(query)[:self.base_compressor.top_n]
        return compressed

safe_retriever = SafeCompressionRetriever(
    base_compressor=cohere_reranker,
    base_retriever=retriever
)

For long documents that exceed the reranker’s token limit (Cohere: 4096 tokens per doc; local models: typically 512), pre-truncate in the compressor:

from langchain.retrievers.document_compressors import DocumentCompressorPipeline
from langchain_community.document_transformers import LongContextReorder

# Truncate each chunk to 512 tokens before reranking
from langchain_text_splitters import TokenTextSplitter
truncator = TokenTextSplitter(chunk_size=512, chunk_overlap=0)

def truncate_docs(docs):
    out = []
    for d in docs:
        chunks = truncator.split_text(d.page_content)
        out.append(Document(page_content=chunks[0], metadata=d.metadata))
    return out

# Pipeline: truncate -> rerank
pipeline_compressor = DocumentCompressorPipeline(
    transformers=[truncate_docs, cohere_reranker]
)

Step 7: Evaluate end-to-end with RAGAS or custom evals

Reranking should improve answer quality, not just retrieval metrics. Run a small eval set through the full chain.

# eval_rag.py
from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy, context_precision
from datasets import Dataset

questions = [
    "What is the refund policy for enterprise customers?",
    "What SLA uptime do you guarantee?",
    "Where do I find the data processing agreement?",
]
ground_truth = [
    "Enterprise customers receive a full refund within 30 days of purchase.",
    "We guarantee 99.9% uptime monthly.",
    "The DPA is available on request from legal@company.com.",
]

# Run both chains
baseline_answers = [rag_chain.invoke(q) for q in questions]
reranked_answers = [rag_chain_cohere.invoke(q) for q in questions]

# Build dataset for RAGAS
def build_dataset(answers):
    return Dataset.from_dict({
        "question": questions,
        "answer": answers,
        "contexts": [[d.page_content for d in retriever.invoke(q)] for q in questions],
        "ground_truth": ground_truth,
    })

baseline_ds = build_dataset(baseline_answers)
reranked_ds = build_dataset(reranked_answers)

print("=== Baseline ===")
print(evaluate(baseline_ds, metrics=[faithfulness, answer_relevancy, context_precision]))

print("\n=== Reranked ===")
print(evaluate(reranked_ds, metrics=[faithfulness, answer_relevancy, context_precision]))

Expect context_precision and answer_relevancy to improve. faithfulness may stay flat — reranking doesn’t fix hallucination, it just feeds better context.

Step 8: Deploy with fallback for provider degradation

If you use Cohere rerank in production, wrap it with a fallback to the local cross-encoder. This pattern mirrors how n4n.ai handles provider fallback at the model layer — same principle applied to the reranking step.

# fallback_rerank.py
from langchain.retrievers.document_compressors import BaseDocumentCompressor
from typing import List
from langchain_core.documents import Document
from langchain_core.callbacks import Callbacks

class FallbackReranker(BaseDocumentCompressor):
    primary: BaseDocumentCompressor
    fallback: BaseDocumentCompressor

    def compress_documents(
        self,
        documents: List[Document],
        query: str,
        callbacks: Callbacks = None,
    ) -> List[Document]:
        try:
            return self.primary.compress_documents(documents, query, callbacks)
        except Exception as e:
            # Log the failure, then fall back
            print(f"Primary reranker failed: {e}. Using fallback.")
            return self.fallback.compress_documents(documents, query, callbacks)

# Local fallback model
local_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-v2-m3")
local_compressor = CrossEncoderReranker(model=local_model, top_n=5)

fallback_compressor = FallbackReranker(primary=cohere_reranker, fallback=local_compressor)

fallback_retriever = ContextualCompressionRetriever(
    base_compressor=fallback_compressor,
    base_retriever=retriever
)

# Use in chain
rag_chain_fallback = (
    {"context": fallback_retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

Test the fallback by temporarily invalidating your Cohere key and verifying the chain still returns answers.

Verification checklist

Before merging to main, confirm:

1. Retrieval recall: On your eval queries, the reranked top-5 contains the ground-truth chunk ≥90% of the time.
2. Latency: P95 rerank + retrieval stays under your SLA (typically <500ms for Cohere, <2s for local GPU).
3. Cost: Cohere rerank calls stay within budget — k * top_n * queries_per_day * $0.001 per 1k docs.
4. Fallback: Kill the primary reranker; the chain degrades gracefully to local model.
5. Observability: Rerank scores appear in logs; alerts fire on score drift.

What to try next

• Hybrid retrieval: Combine BM25 (keyword) + dense vectors before reranking. LangChain’s EnsembleRetriever makes this trivial.
• Query rewriting: Add an LLM step to expand ambiguous queries before retrieval — improves recall for the reranker.
• Reranker fine-tuning: If you have labeled relevance data, fine-tune a cross-encoder on your domain. sentence-transformers supports this with CrossEncoder training loops.
• Caching: Cache reranker scores for repeated queries. A simple Redis lookup on (query_hash, doc_id) saves 80%+ of rerank calls in chat workloads.

Reranking turns a noisy retrieval signal into a clean context window. The code above is production-ready — copy, adapt, and ship.