9.4 Chunking, Reranking & Hybrid Search¶

# fixed_chunking.py — minimal, dependency-free fixed-length chunker
from __future__ import annotations
from typing import Iterator

def fixed_chunk(
    text: str,
    chunk_size: int = 512,
    overlap: int = 64,
    tokenizer=None,      # a callable str->list[int]; falls back to whitespace split
) -> Iterator[str]:
    """
    Yield overlapping token-level chunks of `text`.

    Args:
        text:        The raw document string.
        chunk_size:  Maximum tokens per chunk.
        overlap:     Number of tokens to repeat from the previous chunk.
        tokenizer:   Optional callable returning token ids. When None, we split
                     on whitespace as a proxy (fast for prototyping).

    Yields:
        Decoded string chunks.

    Example:
        >>> chunks = list(fixed_chunk("word " * 1000, chunk_size=10, overlap=2))
        >>> len(chunks)   # ceil((1000-10) / (10-2)) + 1 = 125 chunks
        125
    """
    if tokenizer is None:
        # Whitespace proxy: each "token" is a word
        tokens = text.split()
        decode = lambda ids: " ".join(ids)  # noqa: E731
    else:
        tokens = tokenizer(text)
        decode = tokenizer.decode  # type: ignore[attr-defined]

    step = chunk_size - overlap
    if step <= 0:
        raise ValueError("chunk_size must be strictly greater than overlap")

    start = 0
    while start < len(tokens):
        end = min(start + chunk_size, len(tokens))
        yield decode(tokens[start:end])
        if end == len(tokens):
            break
        start += step


# --- Quick demo ---
if __name__ == "__main__":
    doc = " ".join([f"word{i}" for i in range(200)])
    chunks = list(fixed_chunk(doc, chunk_size=50, overlap=10))
    print(f"Produced {len(chunks)} chunks from 200 tokens")
    print(f"First chunk length: {len(chunks[0].split())} words")
    print(f"Last chunk:         {chunks[-1][:60]}...")

# semantic_chunking.py — embed-based semantic chunker
from __future__ import annotations
import numpy as np
from typing import Callable


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity between two 1-D vectors."""
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))


def semantic_chunk(
    sentences: list[str],
    embed_fn: Callable[[list[str]], np.ndarray],
    threshold: float = 0.75,
    min_sentences: int = 3,
) -> list[str]:
    """
    Group consecutive sentences into chunks based on embedding similarity.

    Args:
        sentences:     Pre-tokenized list of sentence strings.
        embed_fn:      Function mapping list[str] -> np.ndarray of shape (N, D).
        threshold:     Similarity threshold; boundaries placed where sim < threshold.
        min_sentences: Minimum sentences per chunk to avoid micro-chunks.

    Returns:
        List of joined chunk strings.

    The key insight: when cos-sim between sentence[i] and sentence[i+1]
    drops sharply, we've hit a topic boundary. This is much more robust
    than arbitrary token counts for structured prose.
    """
    if not sentences:
        return []

    # Embed all sentences in one batch (efficient for API calls)
    embeddings = embed_fn(sentences)  # shape (N, D)

    # Compute consecutive similarities
    sims = [
        cosine_similarity(embeddings[i], embeddings[i + 1])
        for i in range(len(sentences) - 1)
    ]

    # Find boundary positions (where similarity is low AND gap >= min_sentences)
    chunks: list[str] = []
    current: list[str] = []
    for i, sent in enumerate(sentences):
        current.append(sent)
        if i < len(sims):
            is_boundary = sims[i] < threshold and len(current) >= min_sentences
            if is_boundary:
                chunks.append(" ".join(current))
                current = []

    if current:  # flush remaining sentences
        chunks.append(" ".join(current))

    return chunks


# --- Illustrative usage (replace embed_fn with your actual model) ---
if __name__ == "__main__":
    import random

    # Mock embed function: random vectors — replace with real embeddings
    def mock_embed(texts: list[str]) -> np.ndarray:
        rng = np.random.default_rng(42)
        return rng.standard_normal((len(texts), 384)).astype(np.float32)

    sample = [f"Sentence {i} about {'topic A' if i < 5 else 'topic B'}."
              for i in range(10)]
    result = semantic_chunk(sample, mock_embed, threshold=0.5, min_sentences=2)
    print(f"Produced {len(result)} semantic chunks")

# structural_chunking.py — Markdown-aware recursive splitter
from __future__ import annotations
import re


HEADING_RE = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)


def split_markdown(
    text: str,
    max_tokens: int = 512,
    words_per_token: float = 0.75,  # English rule of thumb: ~1.3 tokens per word
) -> list[dict]:
    """
    Split a Markdown document respecting heading hierarchy.

    Each returned dict has:
      - 'text':     the chunk content
      - 'heading':  the nearest parent heading (for metadata)
      - 'level':    heading depth (1-6; 0 = preamble)

    Strategy:
      1. Split on headings first.
      2. If a section exceeds max_tokens, sub-split on paragraph breaks.
      3. Preserve heading as metadata for downstream metadata filtering.
    """
    # English text runs ~1.3 BPE tokens per whitespace word, so a 512-token
    # budget is ~384 words. (Multiply, do not divide: dividing would let each
    # chunk grow to ~680 words ≈ 900 tokens and silently overflow the encoder.)
    max_words = int(max_tokens * words_per_token)
    chunks: list[dict] = []

    # Find all heading positions
    boundaries = [
        (m.start(), m.group(1), m.group(2).strip()) for m in HEADING_RE.finditer(text)
    ]
    boundaries.append((len(text), "", ""))  # sentinel

    prev_end = 0
    current_heading = ""
    current_level = 0

    for i, (pos, hashes, heading_text) in enumerate(boundaries):
        section = text[prev_end:pos].strip()
        if section:
            words = section.split()
            if len(words) <= max_words:
                chunks.append({
                    "text": section,
                    "heading": current_heading,
                    "level": current_level,
                })
            else:
                # Sub-split on blank lines (paragraphs)
                paragraphs = re.split(r"\n\s*\n", section)
                for para in paragraphs:
                    para = para.strip()
                    if para:
                        chunks.append({
                            "text": para,
                            "heading": current_heading,
                            "level": current_level,
                        })

        current_heading = heading_text
        current_level = len(hashes)
        prev_end = pos

    return chunks

# real_chunkers.py — the production equivalents of §2.1–2.3
# pip install langchain-text-splitters llama-index-core

# --- LangChain: recursive character splitting (the workhorse default) ---
from langchain_text_splitters import (
    RecursiveCharacterTextSplitter,
    MarkdownHeaderTextSplitter,
)

# Tries separators in order — paragraphs, then lines, then sentences, then
# words — so it only cuts mid-sentence when it has no other option.
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512, chunk_overlap=64,
    separators=["\n\n", "\n", ". ", " ", ""],
)
document = "\n\n".join(f"Paragraph {i}. " + "word " * 100 for i in range(20))
chunks = splitter.split_text(document)

# To count in *tokens* rather than characters, bind a real tokenizer:
#   RecursiveCharacterTextSplitter.from_huggingface_tokenizer(tok, chunk_size=512)
#   RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=512)

# --- LangChain: structure-aware (the §2.3 equivalent) ---
# Splits on headings and attaches them as chunk metadata for later filtering.
md_splitter = MarkdownHeaderTextSplitter(
    headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")]
)

# --- LlamaIndex: sentence-window and semantic (the §2.2 equivalent) ---
from llama_index.core.node_parser import SentenceSplitter, SemanticSplitterNodeParser
sentence_parser = SentenceSplitter(chunk_size=512, chunk_overlap=64)
# SemanticSplitterNodeParser(embed_model=..., breakpoint_percentile_threshold=95)
# implements §2.2's embedding-similarity boundary detection, with the threshold
# expressed as a percentile of observed similarity drops rather than an absolute.

# late_chunking.py — illustrative late chunking with a HuggingFace model
from __future__ import annotations
import torch
from transformers import AutoTokenizer, AutoModel
from typing import NamedTuple


class LateChunk(NamedTuple):
    text: str
    embedding: torch.Tensor  # shape (D,)
    start_char: int
    end_char: int


def late_chunk_document(
    document: str,
    chunk_boundaries: list[tuple[int, int]],  # (start_char, end_char) pairs
    model_name: str = "jinaai/jina-embeddings-v2-base-en",
    device: str = "cpu",
) -> list[LateChunk]:
    """
    Apply late chunking: encode full document once, pool over chunk spans.

    Args:
        document:          Full document string.
        chunk_boundaries:  List of (start_char, end_char) defining each chunk.
                           These can be obtained from structural_chunking above,
                           or any other boundary detection method.
        model_name:        HuggingFace encoder model (must support long context).
        device:            'cpu' or 'cuda'.

    Returns:
        List of LateChunk named tuples with text + contextual embedding.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(device).eval()

    # Tokenize entire document, keep char-to-token mapping
    encoding = tokenizer(
        document,
        return_tensors="pt",
        return_offsets_mapping=True,  # crucial: gives (char_start, char_end) per token
        truncation=True,              # truncate to model max length if needed
        max_length=8192,
    )
    offset_mapping = encoding.pop("offset_mapping")[0]  # (seq_len, 2)
    input_ids = encoding["input_ids"].to(device)
    attention_mask = encoding["attention_mask"].to(device)

    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        # token_embeddings: (1, seq_len, hidden_size)
        token_embeddings = outputs.last_hidden_state[0]  # (seq_len, D)

    chunks: list[LateChunk] = []
    for start_char, end_char in chunk_boundaries:
        # Find which token indices correspond to this character span
        token_mask = (
            (offset_mapping[:, 0] >= start_char) &
            (offset_mapping[:, 1] <= end_char)
        )
        span_embeddings = token_embeddings[token_mask]  # (span_len, D)

        if span_embeddings.shape[0] == 0:
            continue  # skip empty spans (e.g., punctuation-only)

        # Mean pooling over the span's token embeddings
        chunk_embedding = span_embeddings.mean(dim=0)  # (D,)

        # L2-normalize for cosine similarity
        chunk_embedding = chunk_embedding / (chunk_embedding.norm() + 1e-9)

        chunks.append(LateChunk(
            text=document[start_char:end_char],
            embedding=chunk_embedding.cpu(),
            start_char=start_char,
            end_char=end_char,
        ))

    return chunks

# real_bm25.py — bm25s: sparse-matrix BM25, same math, ~100x faster
# pip install bm25s PyStemmer
import bm25s
import Stemmer

corpus = [
    "Reciprocal rank fusion combines ranked lists without score normalization.",
    "CVE-2023-44487 is the HTTP/2 Rapid Reset denial-of-service vulnerability.",
    "BM25 weights rare terms heavily through the IDF factor.",
]

stemmer = Stemmer.Stemmer("english")                      # optional but helps recall
corpus_tokens = bm25s.tokenize(corpus, stopwords="en", stemmer=stemmer)

retriever = bm25s.BM25(method="lucene")                   # Lucene's BM25 variant
retriever.index(corpus_tokens)                            # builds the sparse matrix

query_tokens = bm25s.tokenize("HTTP/2 rapid reset CVE", stopwords="en", stemmer=stemmer)
results, scores = retriever.retrieve(query_tokens, k=2)   # shapes: (1, k) each
for rank in range(results.shape[1]):
    print(f"{rank + 1}. score={scores[0, rank]:.3f}  {corpus[results[0, rank]]}")

# hybrid_search.py — BM25 + dense retrieval with RRF fusion
from __future__ import annotations
from dataclasses import dataclass, field
from collections import defaultdict
from typing import Callable
import math
import numpy as np


@dataclass
class Document:
    id: str
    text: str
    embedding: np.ndarray | None = field(default=None, repr=False)


# ─── BM25 implementation ────────────────────────────────────────────────────

class BM25Index:
    """Minimal BM25 index for a list of Documents."""

    def __init__(self, docs: list[Document], k1: float = 1.5, b: float = 0.75):
        self.docs = docs
        self.k1, self.b = k1, b
        self._build(docs)

    def _build(self, docs: list[Document]) -> None:
        # Term frequency per document
        self.tf: list[dict[str, int]] = []
        self.doc_len: list[int] = []
        self.df: dict[str, int] = defaultdict(int)

        for doc in docs:
            tokens = doc.text.lower().split()
            freq: dict[str, int] = defaultdict(int)
            for tok in tokens:
                freq[tok] += 1
            self.tf.append(dict(freq))
            self.doc_len.append(len(tokens))
            for tok in freq:
                self.df[tok] += 1

        self.N = len(docs)
        self.avgdl = sum(self.doc_len) / max(self.N, 1)

    def _idf(self, term: str) -> float:
        n = self.df.get(term, 0)
        return math.log((self.N - n + 0.5) / (n + 0.5) + 1.0)

    def score(self, query: str, doc_idx: int) -> float:
        tokens = query.lower().split()
        dl = self.doc_len[doc_idx]
        score = 0.0
        for tok in tokens:
            f = self.tf[doc_idx].get(tok, 0)
            idf = self._idf(tok)
            denom = f + self.k1 * (1 - self.b + self.b * dl / self.avgdl)
            score += idf * (f * (self.k1 + 1)) / denom
        return score

    def search(self, query: str, top_k: int = 10) -> list[tuple[int, float]]:
        """Return (doc_idx, bm25_score) sorted descending."""
        scores = [(i, self.score(query, i)) for i in range(self.N)]
        scores.sort(key=lambda x: -x[1])
        return scores[:top_k]


# ─── Dense retrieval ─────────────────────────────────────────────────────────

def dense_search(
    query_embedding: np.ndarray,
    doc_embeddings: np.ndarray,  # (N, D)
    top_k: int = 10,
) -> list[tuple[int, float]]:
    """Return (doc_idx, cosine_sim) sorted descending."""
    sims = doc_embeddings @ query_embedding  # dot product = cosine sim if normalized
    idx = np.argsort(-sims)[:top_k]
    return [(int(i), float(sims[i])) for i in idx]


# ─── Reciprocal Rank Fusion ──────────────────────────────────────────────────

def reciprocal_rank_fusion(
    rankings: list[list[tuple[int, float]]],
    k: int = 60,
) -> list[tuple[int, float]]:
    """
    Fuse multiple ranked lists via RRF.

    Args:
        rankings: Each element is a list of (doc_idx, score) sorted by score desc.
                  The actual scores are ignored; only ranks matter.
        k:        RRF smoothing constant (default 60, from Cormack et al. 2009).

    Returns:
        Fused list of (doc_idx, rrf_score) sorted descending.
    """
    rrf_scores: dict[int, float] = defaultdict(float)
    for ranked_list in rankings:
        for rank, (doc_idx, _score) in enumerate(ranked_list, start=1):
            rrf_scores[doc_idx] += 1.0 / (k + rank)

    fused = sorted(rrf_scores.items(), key=lambda x: -x[1])
    return fused


# ─── Putting it together ──────────────────────────────────────────────────────

def hybrid_search(
    query: str,
    query_embedding: np.ndarray,
    bm25_index: BM25Index,
    doc_embeddings: np.ndarray,
    top_k: int = 10,
    rrf_k: int = 60,
) -> list[tuple[Document, float]]:
    """
    Hybrid BM25 + dense search with RRF fusion.

    Returns top_k (Document, rrf_score) pairs sorted by descending fused rank.
    """
    bm25_results = bm25_index.search(query, top_k=top_k * 2)
    dense_results = dense_search(query_embedding, doc_embeddings, top_k=top_k * 2)

    fused = reciprocal_rank_fusion([bm25_results, dense_results], k=rrf_k)
    docs = bm25_index.docs
    return [(docs[idx], score) for idx, score in fused[:top_k]]

# splade_retrieval.py — learned sparse retrieval with sentence-transformers v5+
# pip install "sentence-transformers>=5.0"
from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-v3")          # MLM head + SpladePooling

corpus = [
    "CVE-2023-44487 is the HTTP/2 Rapid Reset denial-of-service vulnerability.",
    "Reciprocal rank fusion combines ranked lists without score normalization.",
]
doc_emb = model.encode_document(corpus)           # sparse tensors, |V|-dimensional
query_emb = model.encode_query("http2 rapid reset attack")

scores = model.similarity(query_emb, doc_emb)     # sparse dot product
print(scores)

# Inspect *why* a document scored: the non-zero dimensions are vocabulary terms.
print(model.decode(doc_emb[0], top_k=8))          # e.g. [('http', 2.1), ('reset', 1.9), ...]

# cross_encoder_rerank.py — cross-encoder reranking with sentence-transformers
from __future__ import annotations
from sentence_transformers import CrossEncoder
import numpy as np


def load_reranker(
    model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
) -> CrossEncoder:
    """
    Load a cross-encoder ONCE, at process startup.

    Common 2026 choices, cheapest first:
      - 'cross-encoder/ms-marco-MiniLM-L-6-v2'   English, 6 layers, ~ms latency
      - 'BAAI/bge-reranker-v2-m3'                multilingual, 568M, strong default
      - 'mixedbread-ai/mxbai-rerank-base-v2'     multilingual, instruction-aware
      - 'Qwen/Qwen3-Reranker-0.6B'               LLM-initialized, prompt-conditioned,
                                                 long context (needs its own API;
                                                 see the note after this listing)
    """
    return CrossEncoder(model_name, max_length=512)


def rerank(
    query: str,
    candidates: list[str],
    model: CrossEncoder,
    top_k: int = 5,
    batch_size: int = 32,
) -> list[tuple[str, float]]:
    """
    Rerank a list of candidate passages using a preloaded cross-encoder.

    Args:
        query:      The user query string.
        candidates: List of passage strings to rerank.
        model:      A CrossEncoder from load_reranker(). Passing the *model*
                    rather than a model name is deliberate: constructing a
                    CrossEncoder pulls weights from disk and moves them to the
                    GPU, which costs hundreds of milliseconds to seconds — far
                    more than the reranking itself. Do it once per process,
                    never once per query.
        top_k:      Number of top passages to return after reranking.
        batch_size: Batch size for model inference.

    Returns:
        List of (passage_text, relevance_score) sorted by descending score.

    Notes:
        - Scores are logits (not probabilities) from the final classification
          head, unless the model card specifies a sigmoid activation. They are
          comparable *within* one query's candidate list and meaningless across
          queries or across models — never threshold on a raw logit without
          calibrating on your own data first.
        - Higher score = more relevant.
    """
    if not candidates:
        return []

    # Build (query, passage) pairs — one per candidate
    pairs = [(query, passage) for passage in candidates]

    # Score all pairs; model handles batching internally
    scores: np.ndarray = model.predict(pairs, batch_size=batch_size)

    # Sort by descending score
    ranked = sorted(zip(candidates, scores), key=lambda x: -x[1])
    return ranked[:top_k]


# --- Illustrative latency numbers ---
# On a single A100:
#   cross-encoder/ms-marco-MiniLM-L-6-v2, 100 candidates, avg passage 200 tokens:
#   ~60ms per query batch
# On CPU (Intel Xeon):
#   ~400ms for the same 100 candidates

# reranker_finetune.py — domain-adapting a cross-encoder (sentence-transformers v4+)
# pip install "sentence-transformers>=4.0" datasets
from datasets import Dataset
from sentence_transformers.cross_encoder import (
    CrossEncoder,
    CrossEncoderTrainer,
    CrossEncoderTrainingArguments,
)
from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss

# Columns are positional: (text_a, text_b, label). Names are free-form.
train_dataset = Dataset.from_dict({
    "query": ["What is RRF?", "What is RRF?"],
    "passage": [
        "Reciprocal rank fusion combines ranked lists using 1/(k+rank).",
        "The capital of France is Paris.",
    ],
    "label": [1.0, 0.0],   # 1 = relevant, 0 = hard negative
})

# num_labels=1 → a single relevance logit (regression/BCE head), not classification.
model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", num_labels=1)
loss = BinaryCrossEntropyLoss(model)

args = CrossEncoderTrainingArguments(
    output_dir="./my-domain-reranker",
    num_train_epochs=1,          # 1 epoch is usually enough for domain adaptation
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,                   # bf16=True on Ampere+ / H100
)

trainer = CrossEncoderTrainer(model=model, args=args,
                              train_dataset=train_dataset, loss=loss)
trainer.train()
model.save_pretrained("./my-domain-reranker/final")

# query_rewriting.py — LLM-powered multi-query expansion
from __future__ import annotations
import json
import re


REWRITE_PROMPT = """\
You are a retrieval expert. Given a user question, produce 3 alternative phrasings
that together cover the semantic space of the question. Return a JSON list of strings.

User question: {question}

Alternative phrasings (JSON list):"""


def expand_query(
    question: str,
    llm_call: callable,  # fn(prompt: str) -> str
    n_variants: int = 3,
) -> list[str]:
    """
    Use an LLM to generate query variants for multi-query retrieval.

    Args:
        question:   Original user question.
        llm_call:   Callable that takes a prompt and returns a completion string.
        n_variants: Number of alternative phrasings requested (soft limit).

    Returns:
        List of query strings including the original.
    """
    prompt = REWRITE_PROMPT.format(question=question)
    response = llm_call(prompt)

    # Extract JSON list from LLM output (handle markdown code fences)
    json_match = re.search(r"\[.*?\]", response, re.DOTALL)
    if json_match:
        try:
            variants = json.loads(json_match.group())
            if isinstance(variants, list):
                return [question] + variants[:n_variants]
        except json.JSONDecodeError:
            pass

    return [question]  # fallback: return original query only

# hyde.py — Hypothetical Document Embeddings for improved dense retrieval
from __future__ import annotations
import numpy as np


HYDE_PROMPT = """\
Write a short, factual paragraph (2-4 sentences) that directly answers the
following question. Do not say "I don't know" — write the best answer you can,
even if uncertain.

Question: {question}

Answer:"""


def hyde_embed(
    question: str,
    llm_call: callable,           # fn(prompt: str) -> str
    embed_fn: callable,           # fn(list[str]) -> np.ndarray
    n_hypotheses: int = 1,
) -> np.ndarray:
    """
    Generate HyDE embedding for a question.

    When n_hypotheses > 1, generate multiple hypothetical documents and
    average their embeddings. This reduces variance from stochastic generation.

    Args:
        question:      User query.
        llm_call:      LLM completion function.
        embed_fn:      Embedding function mapping list[str] -> (N, D) array.
        n_hypotheses:  Number of hypothetical documents to generate.

    Returns:
        Averaged embedding vector of shape (D,), L2-normalized.
    """
    hypotheses = []
    prompt = HYDE_PROMPT.format(question=question)
    for _ in range(n_hypotheses):
        hypothetical_doc = llm_call(prompt)
        hypotheses.append(hypothetical_doc.strip())

    # Embed all hypotheses and average
    embeddings = embed_fn(hypotheses)  # (n_hypotheses, D)
    mean_embedding = embeddings.mean(axis=0)

    # L2-normalize so downstream cosine search still works
    norm = np.linalg.norm(mean_embedding)
    return mean_embedding / (norm + 1e-9)


# --- When to use HyDE vs raw query embedding ---
#
# Use HyDE when:
#  - Your corpus uses formal or technical language that differs from query style
#  - Queries are short and ambiguous (e.g., "transformer memory management")
#  - You can afford 1-2 extra LLM calls per query (latency budget ~200ms)
#
# Avoid HyDE when:
#  - Queries are already long and specific
#  - LLM hallucinations in the hypothesis could mislead retrieval
#  - Low-latency SLA (< 50ms) prevents extra LLM calls

# metadata_filtering.py — metadata-aware retrieval with Qdrant
from __future__ import annotations
from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue, Range
import numpy as np


def filtered_search(
    client: QdrantClient,
    collection_name: str,
    query_vector: np.ndarray,
    filters: dict,              # simplified filter spec (see below)
    top_k: int = 10,
) -> list[dict]:
    """
    Perform vector search with metadata pre-filtering in Qdrant.

    Args:
        client:          Qdrant client connected to your instance.
        collection_name: Name of the Qdrant collection.
        query_vector:    Query embedding, shape (D,).
        filters:         Dict with optional keys:
                         - 'category': str (exact match)
                         - 'date_after': int (Unix timestamp)
                         - 'date_before': int (Unix timestamp)
        top_k:           Number of results to return.

    Returns:
        List of payload dicts for matched documents.

    Qdrant supports filtering directly on HNSW traversal, making pre-filtering
    efficient even on large corpora (millions of documents).
    """
    conditions = []

    if "category" in filters:
        conditions.append(
            FieldCondition(key="category", match=MatchValue(value=filters["category"]))
        )

    if "date_after" in filters or "date_before" in filters:
        conditions.append(
            FieldCondition(
                key="timestamp",
                range=Range(
                    gte=filters.get("date_after"),
                    lte=filters.get("date_before"),
                ),
            )
        )

    qdrant_filter = Filter(must=conditions) if conditions else None

    # `query_points` is the current unified Query API (qdrant-client >= 1.10);
    # the older `client.search(...)` is deprecated but still functional. The
    # same endpoint also serves hybrid retrieval: pass `prefetch=[...]` with a
    # dense and a sparse sub-query plus `query=FusionQuery(fusion=Fusion.RRF)`
    # to run Section 3.1's fusion server-side.
    response = client.query_points(
        collection_name=collection_name,
        query=query_vector.tolist(),
        query_filter=qdrant_filter,
        limit=top_k,
        with_payload=True,
    )

    return [hit.payload for hit in response.points]

# parent_child_retrieval.py — parent document retrieval pattern
from __future__ import annotations
from dataclasses import dataclass
from typing import Any


@dataclass
class ChildChunk:
    id: str
    parent_id: str   # foreign key to parent document
    text: str        # small chunk used for embedding/retrieval
    embedding: Any   # np.ndarray


@dataclass
class ParentDocument:
    id: str
    text: str        # larger context returned to the LLM
    metadata: dict


def parent_child_retrieve(
    query_embedding,
    child_index,           # your vector index over ChildChunk objects
    parent_store: dict,    # parent_id -> ParentDocument
    top_k: int = 5,
) -> list[ParentDocument]:
    """
    Retrieve top-k child chunks, then return their parent documents.

    Deduplicates parents so the same parent is not returned twice even if
    multiple of its children ranked highly.
    """
    # Retrieve top-k*3 children to ensure we have enough unique parents
    child_hits = child_index.search(query_embedding, top_k=top_k * 3)

    seen_parent_ids: set[str] = set()
    parents: list[ParentDocument] = []

    for child in child_hits:
        pid = child.parent_id
        if pid not in seen_parent_ids:
            seen_parent_ids.add(pid)
            if pid in parent_store:
                parents.append(parent_store[pid])
        if len(parents) >= top_k:
            break

    return parents

# rag_pipeline.py — production-grade RAG pipeline
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Callable
import numpy as np
import logging

logger = logging.getLogger(__name__)


@dataclass
class RAGConfig:
    # Chunking
    chunk_size: int = 512
    chunk_overlap: int = 64
    chunking_strategy: str = "fixed"  # "fixed" | "semantic" | "structural"

    # Retrieval
    bm25_top_k: int = 40          # BM25 candidates before fusion
    dense_top_k: int = 40         # Dense candidates before fusion
    rrf_k: int = 60                # RRF smoothing constant

    # Query expansion
    use_hyde: bool = False          # enable HyDE
    n_query_variants: int = 1       # 1 = no expansion

    # Reranking
    use_reranker: bool = True
    reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"
    rerank_candidates: int = 50     # fused hits actually scored by the cross-encoder
    reranker_top_k: int = 5         # final chunks passed to LLM

    # Metadata filter
    metadata_filter: dict = field(default_factory=dict)


@dataclass
class RAGResult:
    query: str
    retrieved_chunks: list[str]
    rrf_scores: list[float]
    reranker_scores: list[float]
    hyde_hypothesis: str | None = None


def run_rag_pipeline(
    query: str,
    bm25_index,           # BM25Index from Section 3
    dense_index,          # vector store with .search(embedding, top_k) -> list[tuple[int, float]]
    embed_fn: Callable,   # list[str] -> np.ndarray of shape (N, D)
    llm_fn: Callable,     # str -> str (for HyDE / query expansion)
    reranker,             # a preloaded CrossEncoder from load_reranker(), or None
    config: RAGConfig,
) -> RAGResult:
    """
    Full RAG retrieval pipeline: query expansion → hybrid retrieval → rerank.
    """

    # ── Step 1: Query representation ──────────────────────────────────────────
    hyde_hypothesis = None
    if config.use_hyde:
        query_embedding = hyde_embed(query, llm_fn, embed_fn)
        hyde_hypothesis = "(HyDE embedding used)"
        logger.info("Using HyDE embedding for query: %s", query[:60])
    else:
        query_embedding = embed_fn([query])[0]

    # ── Step 2: Hybrid retrieval ───────────────────────────────────────────────
    bm25_results = bm25_index.search(query, top_k=config.bm25_top_k)
    dense_results_raw = dense_index.search(query_embedding, top_k=config.dense_top_k)
    # Normalize dense results to (idx, score) format (indices must match the
    # corpus indices used by bm25_index.docs, same contract as dense_search())
    dense_results = [(int(idx), float(score)) for idx, score in dense_results_raw]

    fused = reciprocal_rank_fusion(
        [bm25_results, dense_results], k=config.rrf_k
    )
    # Map back to document texts
    all_docs = bm25_index.docs
    fused_chunks = [(all_docs[idx].text, score) for idx, score in fused]
    rrf_scores = [s for _, s in fused_chunks]
    candidate_texts = [t for t, _ in fused_chunks]

    logger.info(
        "Hybrid retrieval: %d BM25 + %d dense → %d fused candidates",
        config.bm25_top_k, config.dense_top_k, len(fused_chunks),
    )

    # ── Step 3: Reranking ──────────────────────────────────────────────────────
    if config.use_reranker and reranker is not None:
        # Note we hand `rerank` the already-loaded model. Only the top
        # `rerank_candidates` fused hits are scored: cross-encoder cost is
        # linear in the candidate count, so this cap is the pipeline's main
        # latency knob (see the Exercise 5 arithmetic).
        reranked = rerank(
            query, candidate_texts[:config.rerank_candidates],
            model=reranker,
            top_k=config.reranker_top_k,
        )
        final_chunks = [t for t, _ in reranked]
        reranker_scores = [float(s) for _, s in reranked]
    else:
        final_chunks = candidate_texts[:config.reranker_top_k]
        reranker_scores = rrf_scores[:config.reranker_top_k]

    return RAGResult(
        query=query,
        retrieved_chunks=final_chunks,
        rrf_scores=rrf_scores[:config.reranker_top_k],
        reranker_scores=reranker_scores,
        hyde_hypothesis=hyde_hypothesis,
    )

# rag_diagnostics.py — evaluate each pipeline stage independently
from __future__ import annotations
from typing import NamedTuple


class RetrievalMetrics(NamedTuple):
    recall_at_k: float        # fraction of queries where relevant doc is in top-k
    mrr: float                # mean reciprocal rank
    ndcg_at_k: float          # normalized discounted cumulative gain at k
    latency_p50_ms: float
    latency_p99_ms: float


def compute_recall_at_k(
    relevant_ids: list[set[str]],    # per-query sets of relevant document ids
    retrieved_ids: list[list[str]],  # per-query ranked retrieved document ids
    k: int,
) -> float:
    """
    Recall@k: for each query, was at least one relevant document in top-k?
    Averaged over all queries.
    """
    assert len(relevant_ids) == len(retrieved_ids), "Must align by query"
    hits = sum(
        1 for rel, ret in zip(relevant_ids, retrieved_ids)
        if rel & set(ret[:k])
    )
    return hits / len(relevant_ids)


def compute_mrr(
    relevant_ids: list[set[str]],
    retrieved_ids: list[list[str]],
) -> float:
    """
    Mean Reciprocal Rank: reciprocal of rank of first relevant document.
    """
    rrs = []
    for rel, ret in zip(relevant_ids, retrieved_ids):
        for rank, doc_id in enumerate(ret, start=1):
            if doc_id in rel:
                rrs.append(1.0 / rank)
                break
        else:
            rrs.append(0.0)
    return sum(rrs) / len(rrs)


def compute_ndcg_at_k(
    relevance: list[dict[str, float]],  # per-query {doc_id: graded relevance}
    retrieved_ids: list[list[str]],
    k: int = 10,
) -> float:
    """
    nDCG@k — the metric to report when relevance is *graded* rather than binary.

    DCG@k = sum_{i=1..k} (2^rel_i - 1) / log2(i + 1)

    The gain is exponential in the grade (so a "perfect" hit is worth much more
    than two "partial" ones) and the discount is logarithmic in the rank (so
    moving a good document from position 5 to position 1 matters, but position
    50 to 45 barely does). Dividing by the IDCG — the DCG of the ideal ordering
    — normalizes to [0, 1] so scores are comparable across queries with
    different numbers of relevant documents.

    This is the metric BEIR reports (nDCG@10), and the one to optimize if your
    reranker will be trained with a listwise loss.
    """
    import math

    def dcg(grades: list[float]) -> float:
        return sum((2.0 ** g - 1.0) / math.log2(i + 2) for i, g in enumerate(grades))

    scores = []
    for rel, ret in zip(relevance, retrieved_ids):
        gains = [rel.get(doc_id, 0.0) for doc_id in ret[:k]]
        ideal = sorted(rel.values(), reverse=True)[:k]
        idcg = dcg(ideal)
        scores.append(dcg(gains) / idcg if idcg > 0 else 0.0)
    return sum(scores) / max(len(scores), 1)


# Typical baseline numbers to aim for on a reasonably clean corpus:
#   Recall@10:  > 0.85 (retrieval stage)
#   MRR:        > 0.70
#   After rerank, Recall@5: > 0.80