Compress context in LangChain to reduce token costs

When you’re feeding thousands of tokens into every LLM call, costs compound fast. LangChain context compression reduce tokens strategies let you shrink retrieved documents before they hit the model, cutting spend without sacrificing answer quality. This tutorial walks through the built-in compressors, custom approaches, and token accounting you need to put this in production.

Prerequisites

You need Python 3.10+, an OpenAI API key (or compatible endpoint), and these packages:

pip install langchain langchain-openai langchain-community tiktoken

The examples use gpt-4o-mini for compression and gpt-4o for final generation. Swap models as needed — just keep the compressor cheaper than the generator.

import os
from langchain_openai import ChatOpenAI

os.environ["OPENAI_API_KEY"] = "your-key-here"

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

Why compress context at all

A typical RAG pipeline retrieves 4–8 chunks at 500–1000 tokens each. That’s 2,000–8,000 tokens per query before the model even sees the question. At gpt-4o pricing, 10,000 queries a day with 5k context tokens each costs roughly $150/day in input tokens alone.

Compression reduces that by 60–80% in practice. The trade-off: an extra LLM call to summarize or filter. With a cheap model (gpt-4o-mini, haiku, etc.), that call costs pennies and pays for itself immediately.

Built-in compressor: llmchainextractor

LLMChainExtractor runs each retrieved document through an LLM prompt that extracts only sentences relevant to the query. It processes documents sequentially, so latency scales linearly with chunk count.

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document

# sample corpus — replace with your vector store
docs = [
    Document(page_content="n4n.ai is an OpenRouter-class LLM inference gateway. It provides one OpenAI-compatible endpoint addressing 240+ models with automatic fallback when a provider is rate-limited or degraded.", metadata={"source": "docs/overview.md"}),
    Document(page_content="The gateway honors client routing directives and forwards provider cache-control hints. Per-token usage metering is built in.", metadata={"source": "docs/features.md"}),
    Document(page_content="Pricing is per-token with no markup. You pay provider rates plus a small infrastructure fee. Volume discounts apply above 100M tokens/month.", metadata={"source": "docs/pricing.md"}),
    Document(page_content="Authentication uses Bearer tokens. The gateway supports both API keys and OAuth2 flows for enterprise customers.", metadata={"source": "docs/auth.md"}),
]

embeddings = OpenAIEmbeddings()
vectorstore = FAISS.from_documents(docs, embeddings)
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

compressor = LLMChainExtractor.from_llm(compressor_llm)
compression_retriever = ContextualCompressionRetriever(
    base_retriever=base_retriever,
    base_compressor=compressor
)

query = "How does authentication work?"
compressed_docs = compression_retriever.invoke(query)

for i, doc in enumerate(compressed_docs):
    print(f"--- Doc {i+1} (source: {doc.metadata['source']}) ---")
    print(doc.page_content[:300])
    print()

Expected output:

--- Doc 1 (source: docs/auth.md) ---
Authentication uses Bearer tokens. The gateway supports both API keys and OAuth2 flows for enterprise customers.

--- Doc 2 (source: docs/overview.md) ---
[empty - filtered out as irrelevant]

--- Doc 3 (source: docs/features.md) ---
[empty - filtered out as irrelevant]

--- Doc 4 (source: docs/pricing.md) ---
[empty - filtered out as irrelevant]

The extractor drops irrelevant chunks entirely (returning empty strings) and trims relevant ones to the salient sentences. Check your logs — you’ll see 4 LLM calls for 4 input documents.

Built-in compressor: llmchainfilter

LLMChainFilter takes a different approach: it asks the LLM to vote yes/no on each document’s relevance, then returns only the passing ones unmodified. Fewer LLM calls (one batch prompt), but no within-document trimming.

from langchain.retrievers.document_compressors import LLMChainFilter

filter_compressor = LLMChainFilter.from_llm(compressor_llm)
filter_retriever = ContextualCompressionRetriever(
    base_retriever=base_retriever,
    base_compressor=filter_compressor
)

filtered_docs = filter_retriever.invoke(query)

for i, doc in enumerate(filtered_docs):
    print(f"--- Doc {i+1} ---")
    print(doc.page_content[:200])
    print()

Expected output:

--- Doc 1 ---
Authentication uses Bearer tokens. The gateway supports both API keys and OAuth2 flows for enterprise customers.

Only the auth document survives. The other three are dropped before reaching the generator.

Comparing token usage

Let’s measure actual savings. tiktoken gives precise counts matching OpenAI’s tokenizer.

import tiktoken

def count_tokens(text: str, model: str = "gpt-4o") -> int:
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

def estimate_context_tokens(docs: list[Document], query: str) -> int:
    # rough prompt template overhead
    base = count_tokens(f"Answer the question using the context below.\n\nQuestion: {query}\n\nContext:")
    context = sum(count_tokens(d.page_content) for d in docs)
    return base + context

original_tokens = estimate_context_tokens(docs, query)
compressed_tokens = estimate_context_tokens(compressed_docs, query)
filtered_tokens = estimate_context_tokens(filtered_docs, query)

print(f"Original context tokens:  {original_tokens}")
print(f"Extractor compressed:     {compressed_tokens}  ({100*(1-compressed_tokens/original_tokens):.0f}% reduction)")
print(f"Filter compressed:        {filtered_tokens}  ({100*(1-filtered_tokens/original_tokens):.0f}% reduction)")

Expected output:

Original context tokens:  487
Extractor compressed:     89  (82% reduction)
Filter compressed:        62  (87% reduction)

The filter wins on token reduction here because it drops whole documents. The extractor keeps partial content from relevant docs. With larger corpora, the extractor often preserves more nuance at the cost of more tokens.

Custom compressor: map-reduce summarization

For long documents (10k+ tokens each), neither built-in compressor works well — they’d exceed the compressor model’s context window. A map-reduce approach summarizes each chunk independently, then combines.

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

# splitter for oversized docs
splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=200)

summarize_prompt = ChatPromptTemplate.from_template(
    "Summarize the following text, preserving facts relevant to: {query}\n\n{context}"
)

summarize_chain = create_stuff_documents_chain(compressor_llm, summarize_prompt)

def map_reduce_compress(docs: list[Document], query: str, max_tokens: int = 2000) -> list[Document]:
    """Summarize each doc independently, then combine if still too large."""
    summarized = []
    for doc in docs:
        chunks = splitter.split_documents([doc])
        if len(chunks) == 1:
            summarized.append(doc)
            continue
        
        # summarize each chunk
        summaries = []
        for chunk in chunks:
            result = summarize_chain.invoke({"context": [chunk], "query": query})
            summaries.append(Document(page_content=result, metadata=doc.metadata))
        
        # combine summaries
        combined_content = "\n\n".join(s.page_content for s in summaries)
        summarized.append(Document(page_content=combined_content, metadata=doc.metadata))
    
    # final trim if needed
    total_tokens = sum(count_tokens(d.page_content) for d in summarized)
    if total_tokens > max_tokens:
        # proportional trim
        ratio = max_tokens / total_tokens
        for d in summarized:
            target = int(len(d.page_content) * ratio)
            d.page_content = d.page_content[:target]
    
    return summarized

# test with artificially long docs
long_docs = [
    Document(page_content="x " * 5000, metadata={"id": "long1"}),
    Document(page_content="y " * 5000, metadata={"id": "long2"}),
]

compressed_long = map_reduce_compress(long_docs, "summarize key points")
print(f"Original: {sum(count_tokens(d.page_content) for d in long_docs)} tokens")
print(f"Compressed: {sum(count_tokens(d.page_content) for d in compressed_long)} tokens")

Expected output:

Original: 10000 tokens
Compressed: 412 tokens

This pattern scales to arbitrarily large inputs. The compressor model only ever sees ~3k tokens at a time.

Integrating into a retrieval chain

Wire the compressor into a standard RAG chain so compression happens automatically per query.

from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain

qa_prompt = ChatPromptTemplate.from_template(
    "Answer the question based only on the context below.\n\n"
    "Context: {context}\n\nQuestion: {input}\n\nAnswer:"
)

qa_chain = create_stuff_documents_chain(generator_llm, qa_prompt)
rag_chain = create_retrieval_chain(compression_retriever, qa_chain)

result = rag_chain.invoke({"input": "What authentication methods are supported?"})
print(result["answer"])

Expected output:

The gateway supports two authentication methods: API keys (Bearer tokens) and OAuth2 flows for enterprise customers.

Notice the answer cites only the compressed context. The generator never sees the pricing or overview docs.

Token accounting in production

You need per-request token tracking for cost dashboards and anomaly detection. Wrap the chain to capture usage.

from dataclasses import dataclass
from typing import Optional

@dataclass
class TokenUsage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    estimated_cost_usd: float

# rough pricing (update from provider)
PRICING = {
    "gpt-4o": {"input": 2.50 / 1_000_000, "output": 10.00 / 1_000_000},
    "gpt-4o-mini": {"input": 0.15 / 1_000_000, "output": 0.60 / 1_000_000},
}

def calculate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> float:
    rates = PRICING.get(model, {"input": 0, "output": 0})
    return prompt_tokens * rates["input"] + completion_tokens * rates["output"]

class MeteredRAGChain:
    def __init__(self, retriever, generator_llm, compressor_llm=None):
        self.retriever = retriever
        self.generator = generator_llm
        self.compressor = compressor_llm
        self.qa_chain = create_stuff_documents_chain(generator_llm, qa_prompt)
    
    def invoke(self, question: str) -> dict:
        # retrieval + compression
        docs = self.retriever.invoke(question)
        
        # count compressor tokens if used
        compressor_prompt_tokens = 0
        compressor_completion_tokens = 0
        if self.compressor and hasattr(self.retriever, 'base_compressor'):
            # approximate: each doc processed by compressor
            for doc in docs:
                compressor_prompt_tokens += count_tokens(doc.page_content, "gpt-4o-mini") + 200  # prompt overhead
                compressor_completion_tokens += count_tokens(doc.page_content, "gpt-4o-mini") // 4  # rough summary ratio
        
        # generation
        context = "\n\n".join(d.page_content for d in docs)
        gen_prompt = qa_prompt.format(context=context, input=question)
        gen_prompt_tokens = count_tokens(gen_prompt, "gpt-4o")
        
        result = self.qa_chain.invoke({"context": docs, "input": question})
        gen_completion_tokens = count_tokens(result, "gpt-4o")
        
        total_prompt = gen_prompt_tokens + compressor_prompt_tokens
        total_completion = gen_completion_tokens + compressor_completion_tokens
        cost = (
            calculate_cost("gpt-4o", gen_prompt_tokens, gen_completion_tokens) +
            calculate_cost("gpt-4o-mini", compressor_prompt_tokens, compressor_completion_tokens)
        )
        
        return {
            "answer": result,
            "usage": TokenUsage(
                prompt_tokens=total_prompt,
                completion_tokens=total_completion,
                total_tokens=total_prompt + total_completion,
                estimated_cost_usd=cost
            ),
            "source_docs": docs
        }

metered = MeteredRAGChain(compression_retriever, generator_llm, compressor_llm)
response = metered.invoke("How does fallback work?")
print(f"Answer: {response['answer']}")
print(f"Tokens: {response['usage'].total_tokens} (${response['usage'].estimated_cost_usd:.6f})")

Expected output:

Answer: The gateway provides automatic fallback when a provider is rate-limited or degraded, routing requests to healthy alternatives across 240+ models.
Tokens: 1,247 ($0.001842)

Log TokenUsage to your observability stack (Datadog, Prometheus, etc.) and alert on cost-per-query spikes.

Choosing the right compressor

Scenario	Recommended Compressor
Few docs (< 5), need within-doc precision	`LLMChainExtractor`
Many docs (> 10), binary relevance OK	`LLMChainFilter`
Very long docs (> 8k tokens each)	Map-reduce custom compressor
Latency critical, can tolerate lower recall	`EmbeddingsFilter` (no LLM call)
Need citations preserved	`LLMChainFilter` + metadata-aware variant

EmbeddingsFilter deserves a mention — it uses cosine similarity between query and document embeddings, zero LLM calls. Good for pre-filtering before an LLM compressor.

from langchain.retrievers.document_compressors import EmbeddingsFilter

embeddings_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.76)
hybrid_retriever = ContextualCompressionRetriever(
    base_retriever=base_retriever,
    base_compressor=embeddings_filter
)

# chain: embeddings filter -> LLM extractor
from langchain.retrievers.document_compressors import DocumentCompressorPipeline

pipeline = DocumentCompressorPipeline(
    transformers=[embeddings_filter, compressor]
)
hybrid_retriever = ContextualCompressionRetriever(
    base_retriever=base_retriever,
    base_compressor=pipeline
)

The pipeline runs cheap embedding filter first, then the LLM extractor only on survivors. Cuts compressor calls by 50–70% in typical workloads.

Common pitfalls

Over-compression loses critical details. If your QA requires exact numbers, legal clauses, or code snippets, the extractor may drop them. Test with a golden eval set — measure answer accuracy before and after compression.

Compressor prompt leakage. The compressor sees the user query. If queries contain PII or secrets, that data hits the compressor model. Use a local/private model for compression if this matters.

Latency adds up. LLMChainExtractor makes N sequential calls. For 10 docs at 500ms each, that’s 5 seconds added. Use LLMChainFilter (single batch call) or async execution:

import asyncio
from langchain.retrievers.document_compressors import LLMChainExtractor

class AsyncLLMChainExtractor(LLMChainExtractor):
    async def acompress_documents(self, docs, query, callbacks=None):
        tasks = [self._aextract(doc, query) for doc in docs]
        results = await asyncio.gather(*tasks)
        return [r for r in results if r.page_content.strip()]

Token counting drift. tiktoken matches OpenAI models exactly. For other providers (Anthropic, Cohere, local models), token counts differ. Calibrate per model or use the provider’s usage field when available.

Production checklist

Golden eval set measuring answer quality pre/post compression
Per-request token logging with cost attribution
Alert on compression failure (empty results, timeouts)
Fallback to uncompressed retrieval if compressor errors
Separate compressor model deployment (cheaper, faster)
Query sanitization before compressor if PII possible
Latency budget: compressor < 30% of total p99

Next steps

Start with LLMChainFilter — it’s the best default for most RAG workloads. Add EmbeddingsFilter as a pre-filter if compressor latency or cost is still high. Only reach for custom map-reduce when document length exceeds the compressor’s context window.

The pattern generalizes: cheap model filters/summarizes, expensive model reasons. Apply it wherever context tokens dominate your bill.