Legal document search with LangChain and hybrid retrieval

Legal document search demands precision that pure semantic search often misses. Legal language relies on specific terminology, citations, and exact phrasing where keyword overlap matters as much as conceptual similarity. This tutorial builds a hybrid retrieval system with LangChain that combines dense embeddings for semantic understanding with BM25 for exact term matching — the combination that actually works for contracts, case law, and regulatory filings.

Prerequisites

You need Python 3.10+ and an OpenAI API key (or compatible endpoint). Install the dependencies:

pip install langchain langchain-openai langchain-community \
    faiss-cpu rank-bm25 pypdf python-docx tiktoken

If you’re running against a gateway like n4n.ai, set OPENAI_BASE_URL to your endpoint and OPENAI_API_KEY to your gateway key. The code below works unchanged.

Project structure

legal-search/
├── data/                    # PDF/DOCX/TXT files go here
├── ingest.py                # One-time indexing script
├── search.py                # Query interface
├── hybrid_retriever.py      # Core hybrid retrieval logic
└── config.py                # Tunable parameters

Configuration first

Keep tunable parameters in one place so you can experiment without hunting through code.

# config.py
from dataclasses import dataclass

@dataclass
class Config:
    # Chunking
    chunk_size: int = 1000
    chunk_overlap: int = 150
    
    # Embeddings
    embedding_model: str = "text-embedding-3-small"
    
    # Retrieval
    dense_k: int = 20          # Candidates from vector store
    sparse_k: int = 20         # Candidates from BM25
    final_k: int = 5           # Results after fusion
    
    # Fusion weights (must sum to 1.0)
    dense_weight: float = 0.6
    sparse_weight: float = 0.4
    
    # Reranking (optional, requires cohere or cross-encoder)
    use_reranker: bool = False
    reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"

CONFIG = Config()

Hybrid retriever implementation

The core insight: dense vectors capture “this clause resembles that clause” while BM25 captures “this clause contains the exact phrase ‘force majeure’.” You need both.

# hybrid_retriever.py
from typing import List, Tuple
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore
from langchain_community.retrievers import BM25Retriever
from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    """
    Combines dense vector search with BM25 sparse retrieval using
    reciprocal rank fusion (RRF). RRF is parameter-light and works
    well without training data.
    """
    
    def __init__(
        self,
        vectorstore: VectorStore,
        documents: List[Document],
        dense_k: int = 20,
        sparse_k: int = 20,
        final_k: int = 5,
        dense_weight: float = 0.6,
        sparse_weight: float = 0.4,
        rrf_k: int = 60
    ):
        self.vectorstore = vectorstore
        self.documents = documents
        self.dense_k = dense_k
        self.sparse_k = sparse_k
        self.final_k = final_k
        self.dense_weight = dense_weight
        self.sparse_weight = sparse_weight
        self.rrf_k = rrf_k
        
        # Build BM25 index from document texts
        tokenized_corpus = [doc.page_content.split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized_corpus)
        self._doc_id_to_index = {id(doc): i for i, doc in enumerate(documents)}
    
    def _dense_search(self, query: str) -> List[Tuple[Document, float]]:
        """Return (doc, score) pairs from vector similarity search."""
        results = self.vectorstore.similarity_search_with_score(
            query, k=self.dense_k
        )
        # FAISS returns distance (lower = better), convert to similarity
        return [(doc, 1.0 / (1.0 + score)) for doc, score in results]
    
    def _sparse_search(self, query: str) -> List[Tuple[Document, float]]:
        """Return (doc, score) pairs from BM25."""
        tokenized_query = query.split()
        scores = self.bm25.get_scores(tokenized_query)
        top_indices = np.argsort(scores)[::-1][:self.sparse_k]
        return [(self.documents[i], float(scores[i])) for i in top_indices]
    
    def _reciprocal_rank_fusion(
        self,
        dense_results: List[Tuple[Document, float]],
        sparse_results: List[Tuple[Document, float]]
    ) -> List[Document]:
        """
        RRF: score = sum(weight / (k + rank)) for each result list.
        Rank is 1-indexed position in that list.
        """
        doc_scores = {}
        
        # Dense rankings
        for rank, (doc, _) in enumerate(dense_results, 1):
            doc_id = id(doc)
            doc_scores[doc_id] = doc_scores.get(doc_id, 0) + \
                self.dense_weight / (self.rrf_k + rank)
        
        # Sparse rankings
        for rank, (doc, _) in enumerate(sparse_results, 1):
            doc_id = id(doc)
            doc_scores[doc_id] = doc_scores.get(doc_id, 0) + \
                self.sparse_weight / (self.rrf_k + rank)
        
        # Sort by fused score
        sorted_docs = sorted(
            doc_scores.items(),
            key=lambda x: x[1],
            reverse=True
        )
        
        # Return top-k documents (deduplicated by id)
        seen = set()
        final = []
        for doc_id, _ in sorted_docs:
            if doc_id not in seen:
                seen.add(doc_id)
                # Find the document object
                for doc in self.documents:
                    if id(doc) == doc_id:
                        final.append(doc)
                        break
            if len(final) >= self.final_k:
                break
        
        return final
    
    def get_relevant_documents(self, query: str) -> List[Document]:
        dense_results = self._dense_search(query)
        sparse_results = self._sparse_search(query)
        return self._reciprocal_rank_fusion(dense_results, sparse_results)
    
    # LangChain Runnable interface
    def invoke(self, input: dict) -> List[Document]:
        return self.get_relevant_documents(input["query"])

Document ingestion pipeline

Legal documents need structure-aware chunking. A 50-page contract split at arbitrary token boundaries loses section context. Use a recursive splitter that respects headers, then enrich chunks with metadata.

# ingest.py
import os
import pickle
from pathlib import Path
from langchain_community.document_loaders import (
    PyPDFLoader,
    Docx2txtLoader,
    TextLoader
)
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from config import CONFIG

DATA_DIR = Path("data")
INDEX_DIR = Path("index")
INDEX_DIR.mkdir(exist_ok=True)

def load_documents(data_dir: Path) -> list:
    """Load all supported documents from data directory."""
    docs = []
    for file_path in data_dir.rglob("*"):
        if file_path.suffix.lower() == ".pdf":
            loader = PyPDFLoader(str(file_path))
        elif file_path.suffix.lower() in [".docx", ".doc"]:
            loader = Docx2txtLoader(str(file_path))
        elif file_path.suffix.lower() in [".txt", ".md"]:
            loader = TextLoader(str(file_path), encoding="utf-8")
        else:
            continue
        
        loaded = loader.load()
        # Tag each chunk with source metadata
        for doc in loaded:
            doc.metadata["source_file"] = file_path.name
            doc.metadata["source_path"] = str(file_path)
        docs.extend(loaded)
    return docs

def chunk_documents(documents: list) -> list:
    """Split documents preserving legal structure."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CONFIG.chunk_size,
        chunk_overlap=CONFIG.chunk_overlap,
        separators=[
            "\n\nArticle ", "\n\nSection ", "\n\n§ ",
            "\n\n", "\n", ". ", " ", ""
        ],
        length_function=len,
    )
    return splitter.split_documents(documents)

def build_index(chunks: list) -> FAISS:
    """Create FAISS index with OpenAI embeddings."""
    embeddings = OpenAIEmbeddings(model=CONFIG.embedding_model)
    vectorstore = FAISS.from_documents(chunks, embeddings)
    return vectorstore

def main():
    print("Loading documents...")
    raw_docs = load_documents(DATA_DIR)
    print(f"Loaded {len(raw_docs)} raw documents")
    
    print("Chunking...")
    chunks = chunk_documents(raw_docs)
    print(f"Created {len(chunks)} chunks")
    
    print("Building vector index...")
    vectorstore = build_index(chunks)
    
    print("Saving index and chunks...")
    vectorstore.save_local(str(INDEX_DIR))
    
    # Save chunks separately for BM25 (FAISS doesn't store full docs reliably)
    with open(INDEX_DIR / "chunks.pkl", "wb") as f:
        pickle.dump(chunks, f)
    
    print(f"Index saved to {INDEX_DIR}")

if __name__ == "__main__":
    main()

Run it once:

python ingest.py

Expected output:

Loading documents...
Loaded 12 raw documents
Chunking...
Created 347 chunks
Building vector index...
Saving index and chunks...
Index saved to index

Search interface

Now wire the hybrid retriever into a query loop. This version includes optional cross-encoder reranking for the final pass — useful when precision at k=3 matters more than latency.

# search.py
import pickle
from pathlib import Path
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from hybrid_retriever import HybridRetriever
from config import CONFIG

INDEX_DIR = Path("index")

LEGAL_QA_PROMPT = ChatPromptTemplate.from_template("""You are a legal research assistant. Answer the question using only the provided context from legal documents.

Context:
{context}

Question: {question}

Instructions:
- Cite specific clauses, sections, or page references when possible
- If the context doesn't contain the answer, say so explicitly
- Distinguish between binding holdings and dicta in case law
- Note jurisdiction and date relevance for statutes/regulations

Answer:""")

def format_docs(docs: list) -> str:
    """Format retrieved documents with source metadata for the LLM."""
    formatted = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source_file", "unknown")
        page = doc.metadata.get("page", "unknown")
        formatted.append(
            f"[Source {i}: {source}, page {page}]\n{doc.page_content}"
        )
    return "\n\n---\n\n".join(formatted)

def build_chain():
    # Load vectorstore and chunks
    embeddings = OpenAIEmbeddings(model=CONFIG.embedding_model)
    vectorstore = FAISS.load_local(
        str(INDEX_DIR),
        embeddings,
        allow_dangerous_deserialization=True
    )
    
    with open(INDEX_DIR / "chunks.pkl", "rb") as f:
        chunks = pickle.load(f)
    
    # Initialize hybrid retriever
    retriever = HybridRetriever(
        vectorstore=vectorstore,
        documents=chunks,
        dense_k=CONFIG.dense_k,
        sparse_k=CONFIG.sparse_k,
        final_k=CONFIG.final_k,
        dense_weight=CONFIG.dense_weight,
        sparse_weight=CONFIG.sparse_weight,
    )
    
    # Optional reranker
    if CONFIG.use_reranker:
        from langchain_community.cross_encoders import HuggingFaceCrossEncoder
        from langchain.retrievers.document_compressors import CrossEncoderReranker
        
        cross_encoder = HuggingFaceCrossEncoder(
            model_name=CONFIG.reranker_model
        )
        reranker = CrossEncoderReranker(model=cross_encoder, top_n=CONFIG.final_k)
        
        # Wrap retriever to apply reranking
        class RerankingRetriever:
            def __init__(self, base_retriever, reranker):
                self.base = base_retriever
                self.reranker = reranker
            
            def get_relevant_documents(self, query: str):
                docs = self.base.get_relevant_documents(query)
                return self.reranker.compress_documents(docs, query)
            
            def invoke(self, input: dict):
                return self.get_relevant_documents(input["query"])
        
        retriever = RerankingRetriever(retriever, reranker)
    
    # LLM for answer generation
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    
    # RAG chain
    chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | LEGAL_QA_PROMPT
        | llm
        | StrOutputParser()
    )
    
    return chain, retriever

def main():
    chain, retriever = build_chain()
    
    print("Legal Document Search Ready")
    print("Type 'exit' to quit\n")
    
    while True:
        query = input("Query: ").strip()
        if query.lower() in ("exit", "quit", "q"):
            break
        
        if not query:
            continue
        
        print("\nRetrieving...")
        # Show retrieved sources first
        docs = retriever.invoke({"query": query})
        print(f"Retrieved {len(docs)} chunks:")
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get("source_file", "unknown")
            page = doc.metadata.get("page", "unknown")
            preview = doc.page_content[:120].replace("\n", " ")
            print(f"  [{i}] {source} p.{page}: {preview}...")
        
        print("\nGenerating answer...")
        answer = chain.invoke(query)
        print(f"\nAnswer:\n{answer}\n")
        print("-" * 80)

if __name__ == "__main__":
    main()

Run the search interface:

python search.py

Example session:

Legal Document Search Ready
Type 'exit' to quit

Query: What are the force majeure notification requirements in the supplier agreement?

Retrieving...
Retrieved 5 chunks:
  [1] supplier_agreement_v3.pdf p.12: ...force majeure event, the affected party shall notify the other party in writing within five (5) business days...
  [2] supplier_agreement_v3.pdf p.13: ...notification shall describe the nature of the force majeure event, its expected duration, and the steps being taken...
  [3] msa_template.docx p.8: ...Force Majeure. Either party may terminate this Agreement without liability if performance is prevented for more than ninety (90) consecutive days...
  [4] supplier_agreement_v3.pdf p.11: ...including but not limited to acts of God, war, terrorism, labor disputes, government actions, and internet infrastructure failures...
  [5] nda_standard.pdf p.3: ...This Agreement shall be governed by the laws of the State of Delaware...

Generating answer...

Answer:
Based on the supplier agreement (v3), the force majeure notification requirements are:

1. **Timing**: The affected party must notify the other party in writing within **five (5) business days** of the force majeure event (supplier_agreement_v3.pdf, p.12).

2. **Content requirements**: The notification must describe:
   - The nature of the force majeure event
   - Its expected duration
   - The steps being taken to mitigate or resume performance (supplier_agreement_v3.pdf, p.13)

3. **Qualifying events**: The agreement includes acts of God, war, terrorism, labor disputes, government actions, and internet infrastructure failures (supplier_agreement_v3.pdf, p.11).

Note: The MSA template (p.8) contains a separate termination right if performance is prevented for more than 90 consecutive days, but this appears to be a different agreement.

--------------------------------------------------------------------------------

Why this works for legal search

Dense vectors catch semantically similar clauses — “termination for convenience” matches “termination without cause” even when wording differs. BM25 catches exact legal terms — “force majeure,” “indemnification,” “material adverse change” — that embeddings sometimes blur. RRF fusion needs no training data and handles the score scale mismatch automatically.

The chunking strategy matters. Legal documents have hierarchical structure (Article → Section → Subsection). The recursive splitter with legal-aware separators keeps related provisions together. Metadata preservation (source file, page number) lets the LLM cite precisely.

Tuning levers

Increase sparse_weight for statute-heavy corpora where exact citation matching dominates. Increase dense_weight for case law where conceptual similarity matters more. Enable the cross-encoder reranker when you’ve validated the retrieval quality and need the last 5-10% precision gain.

Production considerations

• Incremental updates: FAISS doesn’t support true incremental adds efficiently. Rebuild nightly or use a vector DB (Pinecone, Weaviate, Qdrant) that does.
• Metadata filtering: Add pre-filtering by jurisdiction, date range, or document type before hybrid retrieval. FAISS supports metadata filtering via similarity_search_with_score kwargs.
• Query rewriting: Legal queries benefit from expansion — “termination clause” → “termination clause OR termination provision OR termination section OR force majeure OR material breach.”
• Evaluation: Build a small gold set (20-50 queries with known relevant passages). Measure recall@k and MRR before trusting the system.

The hybrid approach isn’t magic — it’s the pragmatic recognition that legal language lives in two regimes simultaneously: the semantic and the exact. Build for both.