RAG: Hybrid Search Optimization Notes — TF-IDF, BM25 & RRF

"""
From-scratch implementation of hybrid search: BM25 + Vector Search + RRF.
need OPENAI API key in environment

All scoring formulas follow RAG.org:
  - TF(t,d) = count(t,d) / |d|
  - IDF(q)  = ln((N - n(q) + 0.5) / (n(q) + 0.5)) + 1        [BM25 variant]
  - BM25    = Σ IDF(qi) * f(qi,D)*(k1+1) / (f(qi,D) + k1*(1 - b + b*|D|/avgdl))
  - RRF     = Σ 1/(k + rank(d))
"""

import math
import re
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any

import openai
from qdrant_client import QdrantClient

# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------

def tokenize(text: str) -> list[str]:
    """Lowercase and extract alphanumeric tokens."""
    return re.findall(r"[a-z0-9]+", text.lower())


# ---------------------------------------------------------------------------
# Corpus loader
# ---------------------------------------------------------------------------

@dataclass
class Document:
    id: str
    tokens: list[str]
    payload: dict[str, Any] = field(default_factory=dict)


def load_corpus(
    collection_name: str = "knowledge_base",
    host: str = "localhost",
    port: int = 6333,
) -> list[Document]:
    """Scroll all points from Qdrant and build an in-memory corpus."""
    client = QdrantClient(host=host, port=port)
    documents: list[Document] = []
    offset = None

    while True:
        results, next_offset = client.scroll(
            collection_name=collection_name,
            offset=offset,
            limit=256,
            with_payload=True,
            with_vectors=False,
        )
        for point in results:
            text = (point.payload or {}).get("text", "")
            documents.append(Document(
                id=str(point.id),
                tokens=tokenize(text),
                payload=point.payload or {},
            ))
        if next_offset is None:
            break
        offset = next_offset
     
    return documents


# ---------------------------------------------------------------------------
# BM25 scoring (from scratch)
# ---------------------------------------------------------------------------

def bm25_search(
    query: str,
    corpus: list[Document],
    limit: int = 15,
    k1: float = 1.5,
    b: float = 0.75,
) -> list[tuple[Document, float]]:
    """
    Score every document against the query using BM25.

    BM25(D, Q) = Σ IDF(qi) * f(qi,D)*(k1+1) / (f(qi,D) + k1*(1-b + b*|D|/avgdl))

    IDF(qi) = ln((N - n(qi) + 0.5) / (n(qi) + 0.5)) + 1
    """
    query_tokens = tokenize(query)
    if not query_tokens or not corpus:
        return []
     
    N = len(corpus)
    avgdl = sum(len(doc.tokens) for doc in corpus) / N

    # Document frequency: how many docs contain each term
    df: dict[str, int] = defaultdict(int)
    for doc in corpus:
        seen = set(doc.tokens)
        for token in seen:
            df[token] += 1
         
    # Precompute IDF for query terms
    idf: dict[str, float] = {}
    for qt in query_tokens:
        n_q = df.get(qt, 0)
        idf[qt] = math.log((N - n_q + 0.5) / (n_q + 0.5)) + 1
     
    # Score each document
    scored: list[tuple[Document, float]] = []
    for doc in corpus:
        score = 0.0
        doc_len = len(doc.tokens)
        # Term frequency map for this document
        tf: dict[str, int] = defaultdict(int)
        for t in doc.tokens:
            tf[t] += 1
         
        for qt in query_tokens:
            f_q = tf.get(qt, 0)
            if f_q == 0:
                continue
            numerator = f_q * (k1 + 1)
            denominator = f_q + k1 * (1 - b + b * doc_len / avgdl)
            score += idf[qt] * numerator / denominator
         
        if score > 0:
            scored.append((doc, score))
         
    scored.sort(key=lambda x: x[1], reverse=True)
    return scored[:limit]


# ---------------------------------------------------------------------------
# Vector search (via Qdrant)
# ---------------------------------------------------------------------------

@dataclass
class VectorResult:
    id: str
    score: float
    payload: dict[str, Any] = field(default_factory=dict)


def vector_search(
    query: str,
    collection_name: str = "knowledge_base",
    host: str = "localhost",
    port: int = 6333,
    limit: int = 15,
) -> list[VectorResult]:
    """Generate an OpenAI embedding and query Qdrant for nearest neighbors."""
    embedding = openai.OpenAI().embeddings.create(
        model="text-embedding-3-small", input=query
    ).data[0].embedding

    client = QdrantClient(host=host, port=port)
    response = client.query_points(
        collection_name=collection_name,
        query=embedding,
        limit=limit,
        with_payload=True,
    )
    points = response.points if hasattr(response, "points") else response

    return [
        VectorResult(id=str(p.id), score=p.score, payload=p.payload or {})
        for p in points
    ]


# ---------------------------------------------------------------------------
# RRF fusion
# ---------------------------------------------------------------------------

@dataclass
class SearchResult:
    id: str
    score: float
    payload: dict[str, Any] = field(default_factory=dict)


def rrf_fuse(
    bm25_results: list[tuple[Document, float]],
    vec_results: list[VectorResult],
    limit: int = 5,
    k: int = 60,
) -> list[SearchResult]:
    """
    Reciprocal Rank Fusion: score(d) = Σ 1/(k + rank(d))
    """
    rrf_scores: dict[str, float] = defaultdict(float)
    payloads: dict[str, dict[str, Any]] = {}

    for rank, (doc, _bm25_score) in enumerate(bm25_results, start=1):
        rrf_scores[doc.id] += 1.0 / (k + rank)
        payloads.setdefault(doc.id, doc.payload)
     
    for rank, vr in enumerate(vec_results, start=1):
        rrf_scores[vr.id] += 1.0 / (k + rank)
        payloads.setdefault(vr.id, vr.payload)
     
    sorted_ids = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)[:limit]

    return [
        SearchResult(id=pid, score=score, payload=payloads[pid])
        for pid, score in sorted_ids
    ]


# ---------------------------------------------------------------------------
# Entry point: hybrid search
# ---------------------------------------------------------------------------

def hybrid_search(
    query: str,
    collection_name: str = "knowledge_base",
    host: str = "localhost",
    port: int = 6333,
    limit: int = 5,
    bm25_k1: float = 1.5,
    bm25_b: float = 0.75,
    rrf_k: int = 60,
) -> list[SearchResult]:
    """
    Full hybrid search pipeline:
      1. Load corpus from Qdrant
      2. BM25 search (from scratch)
      3. Vector search (via Qdrant + OpenAI embeddings)
      4. RRF fusion
    """
    prefetch = limit * 3

    # 1. Build BM25 corpus
    corpus = load_corpus(collection_name, host, port)

    # 2. BM25 search
    bm25_results = bm25_search(query, corpus, limit=prefetch, k1=bm25_k1, b=bm25_b)

    # 3. Vector search
    vec_results = vector_search(query, collection_name, host, port, limit=prefetch)

    # 4. RRF fusion
    return rrf_fuse(bm25_results, vec_results, limit=limit, k=rrf_k)


def test_machine_learning_notes_exist():
    """Query 'machine learning notes' and verify relevant results are returned."""
    results = hybrid_search("machine learning notes", limit=5)

    # At least one result should come back
    assert len(results) > 0, "Expected at least one result for 'machine learning notes'"

    # All RRF scores should be positive
    for r in results:
        assert r.score > 0, f"RRF score should be positive, got {r.score}"
     
    # At least one result should mention machine learning in its text payload
    texts = [r.payload.get("text", "").lower() for r in results]
    assert any(
        "machine learning" in t for t in texts
    ), f"No result contained 'machine learning'. Top texts: {[t[:80] for t in texts]}"

    # Print results for visibility
    for i, r in enumerate(results, 1):
        snippet = r.payload.get("text", "")[:120].replace("\n", " ")
        print(f"  #{i}  rrf_score={r.score:.5f}  id={r.id}")
        print(f"       source={r.payload.get('source_file', 'N/A')}")
        print(f"       text={snippet}...")
     
     
def main():
    print("Hello from rrf_from_scratch!")
    test_machine_learning_notes_exist()
    print("Test passed!")


if __name__ == "__main__":
    main()