Hybrid Search RAG That Actually Works: BM25 + Vectors + Reranking in Python

pip install rank-bm25 sentence-transformers numpy

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import List, Dict, Tuple
import numpy as np

from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer, CrossEncoder

# =========================
# Data contracts
# =========================

@dataclass(frozen=True)
class DocChunk:
    chunk_id: str
    text: str
    meta: Dict[str, str]

@dataclass(frozen=True)
class SearchHit:
    chunk: DocChunk
    score: float
    source: str  # "bm25" | "vector" | "hybrid" | "rerank"

# =========================
# Utilities
# =========================

def normalize_text(s: str) -> str:
    s = s.strip()
    s = re.sub(r"\s+", " ", s)
    return s

def tokenize_for_bm25(text: str) -> List[str]:
    # Simple tokenization: lowercase + alphanumeric words
    text = text.lower()
    return re.findall(r"[a-z0-9]+", text)

def cosine_sim_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    # a: (n, d), b: (m, d)
    a_norm = a / (np.linalg.norm(a, axis=1, keepdims=True) + 1e-12)
    b_norm = b / (np.linalg.norm(b, axis=1, keepdims=True) + 1e-12)
    return a_norm @ b_norm.T

# =========================
# Hybrid Search RAG Engine
# =========================

class HybridSearchRAG:
    """
    Hybrid Search RAG:
    - BM25 for lexical retrieval
    - Vector search for semantic retrieval
    - Merge & score normalization
    - Cross-encoder reranking for precision
    """

    def __init__(
        self,
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        rerank_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
    ) -> None:
        self.embedder = SentenceTransformer(embedding_model)
        self.reranker = CrossEncoder(rerank_model)

        self.chunks: List[DocChunk] = []
        self._bm25: BM25Okapi | None = None
        self._bm25_tokens: List[List[str]] = []
        self._embeddings: np.ndarray | None = None

    # ---------
    # Indexing
    # ---------

    def index(self, chunks: List[DocChunk]) -> None:
        self.chunks = [DocChunk(c.chunk_id, normalize_text(c.text), c.meta) for c in chunks]

        # BM25 index
        self._bm25_tokens = [tokenize_for_bm25(c.text) for c in self.chunks]
        self._bm25 = BM25Okapi(self._bm25_tokens)

        # Vector index
        texts = [c.text for c in self.chunks]
        self._embeddings = np.array(self.embedder.encode(texts, normalize_embeddings=True))

    # ------------
    # Retrieval
    # ------------

    def bm25_search(self, query: str, top_n: int = 12) -> List[SearchHit]:
        if not self._bm25:
            raise RuntimeError("BM25 index not built. Call index() first.")

        q_tokens = tokenize_for_bm25(query)
        scores = self._bm25.get_scores(q_tokens)
        idxs = np.argsort(scores)[::-1][:top_n]

        hits = []
        for i in idxs:
            hits.append(SearchHit(self.chunks[int(i)], float(scores[int(i)]), "bm25"))
        return hits

    def vector_search(self, query: str, top_n: int = 12) -> List[SearchHit]:
        if self._embeddings is None:
            raise RuntimeError("Vector index not built. Call index() first.")

        q_emb = np.array(self.embedder.encode([query], normalize_embeddings=True))
        sims = cosine_sim_matrix(q_emb, self._embeddings)[0]  # (num_chunks,)
        idxs = np.argsort(sims)[::-1][:top_n]

        hits = []
        for i in idxs:
            hits.append(SearchHit(self.chunks[int(i)], float(sims[int(i)]), "vector"))
        return hits

    # ------------
    # Hybrid merge
    # ------------

    @staticmethod
    def _minmax_norm(scores: List[float]) -> List[float]:
        if not scores:
            return scores
        mn, mx = min(scores), max(scores)
        if abs(mx - mn) < 1e-12:
            return [1.0 for _ in scores]
        return [(s - mn) / (mx - mn) for s in scores]

    def hybrid_search(
        self,
        query: str,
        bm25_top_n: int = 12,
        vec_top_n: int = 12,
        hybrid_top_m: int = 16,
        bm25_weight: float = 0.45,
        vec_weight: float = 0.55,
    ) -> List[SearchHit]:
        """
        Returns merged candidates with normalized weighted scores.
        """
        bm25_hits = self.bm25_search(query, top_n=bm25_top_n)
        vec_hits = self.vector_search(query, top_n=vec_top_n)

        # Normalize within each channel
        bm25_norm = self._minmax_norm([h.score for h in bm25_hits])
        vec_norm = self._minmax_norm([h.score for h in vec_hits])

        merged: Dict[str, Tuple[DocChunk, float]] = {}

        for h, ns in zip(bm25_hits, bm25_norm):
            merged[h.chunk.chunk_id] = (h.chunk, merged.get(h.chunk.chunk_id, (h.chunk, 0.0))[1] + ns * bm25_weight)

        for h, ns in zip(vec_hits, vec_norm):
            merged[h.chunk.chunk_id] = (h.chunk, merged.get(h.chunk.chunk_id, (h.chunk, 0.0))[1] + ns * vec_weight)

        # Sort merged candidates
        merged_hits = [
            SearchHit(chunk=c, score=s, source="hybrid")
            for (c, s) in merged.values()
        ]
        merged_hits.sort(key=lambda x: x.score, reverse=True)

        return merged_hits[:hybrid_top_m]

    # ------------
    # Reranking
    # ------------

    def rerank(self, query: str, candidates: List[SearchHit], top_k: int = 6) -> List[SearchHit]:
        """
        Cross-encoder reranking: (query, chunk_text) -> relevance score
        """
        pairs = [(query, h.chunk.text) for h in candidates]
        rerank_scores = self.reranker.predict(pairs)

        reranked = []
        for h, rs in zip(candidates, rerank_scores):
            reranked.append(SearchHit(h.chunk, float(rs), "rerank"))

        reranked.sort(key=lambda x: x.score, reverse=True)
        return reranked[:top_k]

    # ------------
    # Context
    # ------------

    def build_context(self, hits: List[SearchHit], max_chars: int = 4500) -> str:
        """
        Builds an LLM-ready context string.
        Keeps it simple: top chunks concatenated with separators.
        """
        parts = []
        total = 0

        for h in hits:
            header = f"[{h.chunk.meta.get('source','doc')} | {h.chunk.chunk_id}]"
            block = f"{header}\n{h.chunk.text}\n"
            if total + len(block) > max_chars:
                break
            parts.append(block)
            total += len(block)

        return "\n---\n".join(parts)

    # ------------
    # End-to-end
    # ------------

    def retrieve(
        self,
        query: str,
        bm25_top_n: int = 12,
        vec_top_n: int = 12,
        hybrid_top_m: int = 16,
        rerank_top_k: int = 6,
    ) -> List[SearchHit]:
        candidates = self.hybrid_search(
            query=query,
            bm25_top_n=bm25_top_n,
            vec_top_n=vec_top_n,
            hybrid_top_m=hybrid_top_m,
        )
        return self.rerank(query, candidates, top_k=rerank_top_k)

# =========================
# Example usage (drop-in)
# =========================

if __name__ == "__main__":
    # Replace with your own chunks from PDFs/web/pages/db
    chunks = [
        DocChunk("c1", "Reset password: go to Settings > Security > Reset.", {"source": "help_center"}),
        DocChunk("c2", "Error E1127 occurs when token is expired. Refresh the token.", {"source": "runbook"}),
        DocChunk("c3", "To terminate plan, contact billing at least 7 days before renewal.", {"source": "policy"}),
        DocChunk("c4", "Refunds are only issued within 14 days of purchase for annual plans.", {"source": "policy"}),
    ]

    rag = HybridSearchRAG()
    rag.index(chunks)

    query = "How do I fix E1127 and why is it happening?"
    hits = rag.retrieve(query)

    print("\nTop reranked hits:\n")
    for h in hits:
        print(f"- {h.score:.4f} | {h.chunk.chunk_id} | {h.chunk.meta.get('source')}")

    context = rag.build_context(hits)
    print("\n\nLLM Context:\n")
    print(context)