Hybrid Search and Reranking: BM25, Embeddings, RRF and an LLM Reranker

"""Hybrid search for the Hearthly help centre. Every search function takes a question and returns passage ids
("p01", "p02", ...), best first."""
import re

import numpy as np
from rank_bm25 import BM25Okapi

import rag

KB = rag.load_kb()
IDS = [p["id"] for p in KB]
BY_ID = {p["id"]: p for p in KB}
TEXTS = [rag.passage_text(p) for p in KB]


# ---------- Step 1: keyword search ----------
def tokenize(text):
    """Lowercase runs of letters and digits, punctuation dropped: "Error E41, HT-2000!" -> ["error", "e41", "ht", "2000"]."""
    # TODO (Step 1): lowercase the text and return its runs of letters and digits, e.g. re.findall(r"[a-z0-9]+", ...)
    raise NotImplementedError("Step 1: write tokenize()")


_bm25 = None


def bm25_search(question, k=5):
    """The k passage ids BM25 scores highest for the question."""
    # TODO (Step 1): build BM25Okapi once from [tokenize(t) for t in TEXTS] (keep it in the global _bm25),
    # score the tokenized question with get_scores(), and return the ids (IDS[i]) of the k highest scores, best first.
    raise NotImplementedError("Step 1: write bm25_search()")


# ---------- Step 2: dense search ----------
def dense_search(question, k=5):
    """The k passage ids whose embeddings are closest to the question's."""
    raise NotImplementedError("dense_search() arrives in Step 2")


# ---------- Step 3: fuse the two ----------
def rrf(rankings, k=60):
    """Reciprocal rank fusion: each list gives an id 1 / (k + rank), rank starting at 1. Ids sorted by total, best first."""
    raise NotImplementedError("rrf() arrives in Step 3")


def hybrid_search(question, k=5, depth=20):
    """Fuse the top `depth` of BM25 and of dense search with rrf() and keep the first k."""
    raise NotImplementedError("hybrid_search() arrives in Step 3")


# ---------- Step 4: rerank ----------
RERANK_PROMPT = """You rank help-centre passages for a customer question.

Question: {question}

Passages:
{passages}

Which passages answer the question? Reply with the passage numbers only, best answer first, separated by commas, for example: 3, 1, 2"""


def rerank_prompt(question, ids):
    """The prompt for rerank(): each passage numbered [1], [2], ... in the order given."""
    raise NotImplementedError("rerank_prompt() arrives in Step 4")


def parse_ranking(reply, ids):
    """Turn the model's reply ("3, 1") into ids, best first. Numbers outside 1..len(ids) and repeats are ignored;
    every id the reply leaves out is appended in its original order, so nothing is ever lost."""
    raise NotImplementedError("parse_ranking() arrives in Step 4")


def rerank(question, ids):
    """Ask the chat model to order the candidate ids. On any error keep the original order."""
    raise NotImplementedError("rerank() arrives in Step 4")


# ---------- Step 5: the pipeline you ship ----------
def pool_report(questions, sizes, depth=20):
    """For each pool size: how often the right passage is in hybrid_search(question, size, depth) (recall), and the
    average length in characters of the rerank prompt for that pool (avg_prompt_chars)."""
    raise NotImplementedError("pool_report() arrives in Step 5")


SETTINGS = {"depth": 20, "rerank_top": 5}


def answer_passages(question, k=3):
    """The shipped pipeline: hybrid search, rerank the top SETTINGS["rerank_top"], return the first k ids."""
    candidates = hybrid_search(question, k=SETTINGS["rerank_top"], depth=SETTINGS["depth"])
    return rerank(question, candidates)[:k]