10 LangChain Document Compressors to Reduce Context Length

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

# Build base vector store retriever
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(
    collection_name="docs",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 8})

# Wrap with compressor
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
compressor = LLMChainExtractor.from_llm(llm)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever
)

docs = compression_retriever.invoke("What are the side effects of metformin?")
for doc in docs:
    print(f"[{len(doc.page_content)} chars] {doc.page_content[:200]}")

from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_openai import ChatOpenAI
from langchain_core.documents import Document

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
extractor = LLMChainExtractor.from_llm(llm)

doc = Document(
    page_content="""
    Metformin is a first-line medication for type 2 diabetes.
    It works by decreasing glucose production in the liver.
    Common side effects include nausea, vomiting, and diarrhea.
    Rare but serious side effects include lactic acidosis.
    The drug was first synthesized in 1922 by Emil Werner.
    It is on the WHO list of essential medicines.
    """,
    metadata={"source": "medical_db"}
)

query = "What are the side effects of metformin?"
compressed = extractor.compress_documents([doc], query)
print(compressed[0].page_content)
# → "Common side effects include nausea, vomiting, and diarrhea.
#    Rare but serious side effects include lactic acidosis."

from langchain.retrievers.document_compressors import LLMChainFilter

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
filter_compressor = LLMChainFilter.from_llm(llm)

docs = [
    Document(page_content="Metformin reduces liver glucose production.", metadata={"id": 1}),
    Document(page_content="The history of diabetes treatment began in ancient Egypt.", metadata={"id": 2}),
    Document(page_content="Lactic acidosis is a rare but serious risk with metformin.", metadata={"id": 3}),
]

query = "What are the risks of taking metformin?"
filtered = filter_compressor.compress_documents(docs, query)
print(f"Kept {len(filtered)}/{len(docs)} documents")
# → Kept 2/3 documents (drops the history doc)

from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

emb_filter = EmbeddingsFilter(
    embeddings=embeddings,
    similarity_threshold=0.76  # tune this per domain
)

docs = [
    Document(page_content="Metformin controls blood sugar in type 2 diabetes."),
    Document(page_content="Ancient Rome had a thriving olive oil trade."),
    Document(page_content="Side effects of metformin include GI upset."),
]

query = "metformin side effects"
filtered = emb_filter.compress_documents(docs, query)
print(f"Kept {len(filtered)} docs after embedding filter")

from langchain_cohere import CohereRerank
from langchain.retrievers import ContextualCompressionRetriever

reranker = CohereRerank(
    model="rerank-english-v3.0",
    top_n=3,  # keep top 3 after reranking
    cohere_api_key="your-cohere-api-key"
)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker,
    base_retriever=base_retriever  # retrieves k=10 initially
)

results = compression_retriever.invoke("diabetes medication side effects")
for doc in results:
    print(doc.metadata.get("relevance_score", "N/A"), doc.page_content[:100])

# pip install flashrank
from langchain_community.document_compressors import FlashrankRerank

local_reranker = FlashrankRerank(
    model_name="ms-marco-MiniLM-L-12-v2",
    top_n=3
)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=local_reranker,
    base_retriever=base_retriever
)

results = compression_retriever.invoke("metformin dosage for elderly patients")
for doc in results:
    print(doc.page_content[:150])

from langchain.retrievers.document_compressors import DocumentCompressorPipeline
from langchain.text_splitter import CharacterTextSplitter

# Stage 1: re-split long chunks into smaller pieces
splitter = CharacterTextSplitter(
    chunk_size=300,
    chunk_overlap=0,
    separator=". "
)

# Stage 2: embedding similarity filter (fast, cheap)
emb_filter = EmbeddingsFilter(
    embeddings=embeddings,
    similarity_threshold=0.75
)

# Stage 3: LLM extraction (accurate, expensive — runs on fewer docs now)
llm_extractor = LLMChainExtractor.from_llm(
    ChatOpenAI(model="gpt-4o-mini", temperature=0)
)

pipeline = DocumentCompressorPipeline(
    transformers=[splitter, emb_filter, llm_extractor]
)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=pipeline,
    base_retriever=base_retriever
)

from langchain.retrievers.document_compressors import LLMListwiseRerank

listwise_reranker = LLMListwiseRerank.from_llm(
    llm=ChatOpenAI(model="gpt-4o", temperature=0),
    num_docs_to_keep=3
)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=listwise_reranker,
    base_retriever=base_retriever
)

results = compression_retriever.invoke("Compare first-line diabetes medications")
for i, doc in enumerate(results):
    print(f"Rank {i+1}: {doc.page_content[:200]}")

from langchain_core.documents import Document
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
from pydantic import BaseModel, Field
from typing import List, Sequence, Optional
from langchain_core.callbacks import Callbacks

class StructuredCompressor(BaseDocumentCompressor):
    llm: object

    class Config:
        arbitrary_types_allowed = True

    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None
    ) -> List[Document]:

        prompt = ChatPromptTemplate.from_messages([
            ("system", "Extract the most relevant excerpt from this document for the query. Return JSON with keys: excerpt, relevance_score (0-1), reasoning."),
            ("human", "Query: {query}\n\nDocument: {document}")
        ])

        chain = prompt | self.llm | JsonOutputParser()

        compressed = []
        for doc in documents:
            try:
                result = chain.invoke({
                    "query": query,
                    "document": doc.page_content
                })
                if result["relevance_score"] > 0.5:
                    compressed.append(Document(
                        page_content=result["excerpt"],
                        metadata={
                            **doc.metadata,
                            "relevance_score": result["relevance_score"],
                            "reasoning": result["reasoning"]
                        }
                    ))
            except Exception:
                pass  # drop malformed responses

        return sorted(
            compressed,
            key=lambda x: x.metadata["relevance_score"],
            reverse=True
        )

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

compression_retriever = ContextualCompressionRetriever(
    base_compressor=custom_compressor,
    base_retriever=base_retriever
)

from langchain_experimental.text_splitter import SemanticChunker
from langchain_openai import OpenAIEmbeddings
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.retrievers.document_compressors import DocumentCompressorPipeline

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# Semantic splitter creates chunks at natural topic boundaries
semantic_splitter = SemanticChunker(
    embeddings=embeddings,
    breakpoint_threshold_type="percentile",
    breakpoint_threshold_amount=90
)

# Then filter by similarity
emb_filter = EmbeddingsFilter(
    embeddings=embeddings,
    similarity_threshold=0.78
)

pipeline = DocumentCompressorPipeline(
    transformers=[semantic_splitter, emb_filter]
)

compression_retriever = ContextualCompressionRetriever(
    base_compressor=pipeline,
    base_retriever=base_retriever
)

from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import InMemoryStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader

# Child splitter for embedding (small, precise)
child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

# Parent splitter for LLM context (larger, richer)
parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000)

store = InMemoryStore()
parent_retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter
)

# Add documents
loader = WebBaseLoader("https://example.com/medical-docs")
docs = loader.load()
parent_retriever.add_documents(docs)

# Combine with LLM compressor for final context reduction
llm_filter = LLMChainFilter.from_llm(ChatOpenAI(model="gpt-4o-mini"))

final_retriever = ContextualCompressionRetriever(
    base_compressor=llm_filter,
    base_retriever=parent_retriever
)

results = final_retriever.invoke("What dosage adjustments are needed for renal impairment?")

# pip install ragas
from ragas import evaluate
from ragas.metrics import context_precision, faithfulness
from datasets import Dataset

questions = ["What are metformin side effects?", "How does metformin work?"]
results_compressed = []

for q in questions:
    compressed_docs = compression_retriever.invoke(q)
    context = "\n\n".join(d.page_content for d in compressed_docs)
    
    from langchain_core.prompts import ChatPromptTemplate
    prompt = ChatPromptTemplate.from_template(
        "Answer using only this context:\n{context}\n\nQuestion: {question}"
    )
    answer = (prompt | ChatOpenAI(model="gpt-4o")).invoke({
        "context": context,
        "question": q
    }).content
    
    results_compressed.append({
        "question": q,
        "answer": answer,
        "contexts": [d.page_content for d in compressed_docs],
        "ground_truth": "Reference answer here"
    })

eval_results = evaluate(
    Dataset.from_list(results_compressed),
    metrics=[context_precision, faithfulness]
)
print(eval_results)

import numpy as np

def find_optimal_threshold(queries, docs, ground_truth_relevant, embeddings):
    thresholds = np.arange(0.60, 0.95, 0.05)
    results = []

    for threshold in thresholds:
        emb_filter = EmbeddingsFilter(
            embeddings=embeddings,
            similarity_threshold=float(threshold)
        )

        total_tp = total_fp = total_fn = 0
        for query, doc_list, relevant_ids in zip(queries, docs, ground_truth_relevant):
            filtered = emb_filter.compress_documents(doc_list, query)
            filtered_ids = {d.metadata["id"] for d in filtered}

            tp = len(filtered_ids & relevant_ids)
            fp = len(filtered_ids - relevant_ids)
            fn = len(relevant_ids - filtered_ids)

            total_tp += tp
            total_fp += fp
            total_fn += fn

        precision = total_tp / (total_tp + total_fp + 1e-9)
        recall = total_tp / (total_tp + total_fn + 1e-9)
        f1 = 2 * precision * recall / (precision + recall + 1e-9)

        results.append({
            "threshold": threshold,
            "precision": precision,
            "recall": recall,
            "f1": f1
        })

    best = max(results, key=lambda x: x["f1"])
    print(f"Optimal threshold: {best['threshold']:.2f} (F1={best['f1']:.3f})")
    return best["threshold"]

import asyncio
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages([
    ("system", "Answer based only on the provided context. Say so if the answer isn't there."),
    ("human", "Context:\n{context}\n\nQuestion: {question}")
])

def format_docs(docs):
    return "\n\n---\n\n".join(
        f"[Source: {doc.metadata.get('source', 'unknown')}]\n{doc.page_content}"
        for doc in docs
    )

rag_chain = (
    RunnableParallel({
        "context": compression_retriever | format_docs,
        "question": RunnablePassthrough()
    })
    | prompt
    | ChatOpenAI(model="gpt-4o", temperature=0)
    | StrOutputParser()
)

# Async for concurrent requests
async def answer_questions(questions: list[str]) -> list[str]:
    tasks = [rag_chain.ainvoke(q) for q in questions]
    return await asyncio.gather(*tasks)

# Stream individual responses
async def stream_answer(question: str):
    async for chunk in rag_chain.astream(question):
        print(chunk, end="", flush=True)