ripvec_core/

hybrid.rs

//! Hybrid semantic + keyword search with Reciprocal Rank Fusion (RRF).
//!
//! [`HybridIndex`] wraps a [`SearchIndex`] (dense vector search) and a
//! [`Bm25Index`] (BM25 keyword search) and fuses their ranked results via
//! Reciprocal Rank Fusion so that chunks appearing high in either list
//! bubble to the top of the combined ranking.

use std::collections::HashMap;
use std::fmt;
use std::str::FromStr;

use crate::bm25::Bm25Index;
use crate::chunk::CodeChunk;
use crate::index::SearchIndex;

/// Controls which retrieval strategy is used during search.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum SearchMode {
    /// Fuse semantic (vector) and keyword (BM25) results via RRF.
    #[default]
    Hybrid,
    /// Dense vector cosine-similarity ranking only.
    Semantic,
    /// BM25 keyword ranking only.
    Keyword,
}

impl fmt::Display for SearchMode {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Hybrid => f.write_str("hybrid"),
            Self::Semantic => f.write_str("semantic"),
            Self::Keyword => f.write_str("keyword"),
        }
    }
}

/// Error returned when a `SearchMode` string cannot be parsed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParseSearchModeError(String);

impl fmt::Display for ParseSearchModeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "unknown search mode {:?}; expected hybrid, semantic, or keyword",
            self.0
        )
    }
}

impl std::error::Error for ParseSearchModeError {}

impl FromStr for SearchMode {
    type Err = ParseSearchModeError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "hybrid" => Ok(Self::Hybrid),
            "semantic" => Ok(Self::Semantic),
            "keyword" => Ok(Self::Keyword),
            other => Err(ParseSearchModeError(other.to_string())),
        }
    }
}

/// Combined semantic + keyword search index with RRF fusion.
///
/// Build once from chunks and pre-computed embeddings; query repeatedly
/// via [`search`](Self::search).
pub struct HybridIndex {
    /// Semantic (dense vector) search index.
    pub semantic: SearchIndex,
    /// BM25 keyword search index.
    bm25: Bm25Index,
}

impl HybridIndex {
    /// Build a `HybridIndex` from raw chunks and their pre-computed embeddings.
    ///
    /// Constructs both the [`SearchIndex`] and [`Bm25Index`] in one call.
    /// `cascade_dim` is forwarded to [`SearchIndex::new`] for optional MRL
    /// cascade pre-filtering.
    ///
    /// # Errors
    ///
    /// Returns an error if the BM25 index cannot be built (e.g., tantivy
    /// schema or writer failure).
    pub fn new(
        chunks: Vec<CodeChunk>,
        embeddings: &[Vec<f32>],
        cascade_dim: Option<usize>,
    ) -> crate::Result<Self> {
        let bm25 = Bm25Index::build(&chunks)?;
        let semantic = SearchIndex::new(chunks, embeddings, cascade_dim);
        Ok(Self { semantic, bm25 })
    }

    /// Assemble a `HybridIndex` from pre-built components.
    ///
    /// Useful when the caller has already constructed the sub-indices
    /// separately (e.g., loaded from a cache).
    #[must_use]
    pub fn from_parts(semantic: SearchIndex, bm25: Bm25Index) -> Self {
        Self { semantic, bm25 }
    }

    /// Search the index and return `(chunk_index, score)` pairs.
    ///
    /// Dispatches based on `mode`:
    /// - [`SearchMode::Semantic`] — pure dense vector search via
    ///   [`SearchIndex::rank`].
    /// - [`SearchMode::Keyword`] — pure BM25 keyword search, truncated to
    ///   `top_k`.
    /// - [`SearchMode::Hybrid`] — retrieves both ranked lists, fuses them
    ///   with [`rrf_fuse`], then truncates to `top_k`.
    ///
    /// Scores are min-max normalized to `[0, 1]` regardless of mode, so
    /// a threshold of 0.5 always means "above midpoint of the score range"
    /// whether the underlying scores are cosine similarity, BM25, or RRF.
    #[must_use]
    pub fn search(
        &self,
        query_embedding: &[f32],
        query_text: &str,
        top_k: usize,
        threshold: f32,
        mode: SearchMode,
    ) -> Vec<(usize, f32)> {
        let mut raw = match mode {
            SearchMode::Semantic => {
                // Fetch more than top_k so normalization has a meaningful range.
                self.semantic
                    .rank_turboquant(query_embedding, top_k.max(100), 0.0)
            }
            SearchMode::Keyword => self.bm25.search(query_text, top_k.max(100)),
            SearchMode::Hybrid => {
                let sem = self
                    .semantic
                    .rank_turboquant(query_embedding, top_k.max(100), 0.0);
                let kw = self.bm25.search(query_text, top_k.max(100));
                rrf_fuse(&sem, &kw, 60.0)
            }
        };

        // Min-max normalize scores to [0, 1] so threshold is model-agnostic.
        if let (Some(max), Some(min)) = (raw.first().map(|(_, s)| *s), raw.last().map(|(_, s)| *s))
        {
            let range = max - min;
            if range > f32::EPSILON {
                for (_, score) in &mut raw {
                    *score = (*score - min) / range;
                }
            } else {
                // All scores identical — normalize to 1.0
                for (_, score) in &mut raw {
                    *score = 1.0;
                }
            }
        }

        // Apply threshold on normalized scores, then truncate
        raw.retain(|(_, score)| *score >= threshold);
        raw.truncate(top_k);
        raw
    }

    /// All chunks in the index.
    #[must_use]
    pub fn chunks(&self) -> &[CodeChunk] {
        &self.semantic.chunks
    }
}

impl crate::searchable::SearchableIndex for HybridIndex {
    fn chunks(&self) -> &[CodeChunk] {
        HybridIndex::chunks(self)
    }

    fn search(&self, query_text: &str, top_k: usize, mode: SearchMode) -> Vec<(usize, f32)> {
        // Trait-shape search: no caller-supplied query embedding.
        // BM25 handles text directly. Semantic and hybrid modes would
        // require an embedded query vector — without one they would
        // search against a zero vector, which matches nothing
        // useful — so we route all three modes through keyword. The
        // caller wanting semantic results should use
        // `search_from_chunk` (the canonical goto-definition pattern)
        // which supplies the source chunk's embedding.
        let _ = mode;
        HybridIndex::search(self, &[], query_text, top_k, 0.0, SearchMode::Keyword)
    }

    fn search_from_chunk(
        &self,
        chunk_idx: usize,
        query_text: &str,
        top_k: usize,
        mode: SearchMode,
    ) -> Vec<(usize, f32)> {
        let embedding = self.semantic.embedding(chunk_idx).unwrap_or_default();
        let effective_mode = if embedding.is_empty() {
            SearchMode::Keyword
        } else {
            mode
        };
        HybridIndex::search(self, &embedding, query_text, top_k, 0.0, effective_mode)
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}

/// Reciprocal Rank Fusion of two ranked lists.
///
/// Each entry in `semantic` and `bm25` is `(chunk_index, _score)`.
/// The fused score for a chunk is the sum of `1 / (k + rank + 1)` across
/// every list the chunk appears in, where `rank` is 0-based.
///
/// Returns all chunks that appear in either list, sorted descending by
/// fused RRF score.
///
/// `k` should typically be 60.0 — a conventional constant that smooths the
/// ranking boost for the very top results.
#[must_use]
pub fn rrf_fuse(semantic: &[(usize, f32)], bm25: &[(usize, f32)], k: f32) -> Vec<(usize, f32)> {
    let mut scores: HashMap<usize, f32> = HashMap::new();

    for (rank, &(idx, _)) in semantic.iter().enumerate() {
        *scores.entry(idx).or_insert(0.0) += 1.0 / (k + rank as f32 + 1.0);
    }
    for (rank, &(idx, _)) in bm25.iter().enumerate() {
        *scores.entry(idx).or_insert(0.0) += 1.0 / (k + rank as f32 + 1.0);
    }

    let mut results: Vec<(usize, f32)> = scores.into_iter().collect();
    results.sort_unstable_by(|a, b| {
        b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)) // stable tie-break by chunk index
    });
    results
}

/// Sigmoid steepness for the PageRank percentile boost. Lower values
/// produce a sharper transition between "below median" (low boost) and
/// "above median" (full boost).
const PAGERANK_SIGMOID_STEEPNESS: f32 = 0.15;

/// Sigmoid-shaped multiplicative boost factor for a single PageRank
/// **percentile** in the corpus (not the raw rank value).
///
/// Returns the multiplier (so the final score is `dense_score * factor`).
///
/// ```text
/// factor = 1 + alpha * sigmoid((percentile - 0.5) / s)
/// sigmoid(z) = 1 / (1 + exp(-z))
/// ```
///
/// where `s = PAGERANK_SIGMOID_STEEPNESS`.
///
/// ## Why this shape, with examples
///
/// The first attempt used a logarithmic saturation curve on raw rank
/// values. That failed because raw ranks in a top-K result set
/// concentrate in a tiny band (max ≈ 0.028 in Tokio), producing
/// uniformly tiny boosts. The next attempt added a "presence floor"
/// for `rank > 0`, which failed because tests also have tiny-but-
/// positive PR from PageRank's damping term — both impl and test
/// cleared the floor equally.
///
/// Switching the input to **percentile in the corpus** fixes both
/// pathologies. A test with no inbound edges sits in the bottom decile
/// of the PR distribution (percentile ≈ 0.05); a typical
/// implementation file sits above the median. The sigmoid then makes
/// the transition between "below median" (no boost) and "above median"
/// (near-full boost) sharp:
///
/// | percentile | sigmoid | boost (α=0.5) |
/// |------------|---------|---------------|
/// | 0.05 (low test) | 0.04 | 1.02× |
/// | 0.30           | 0.21 | 1.10× |
/// | 0.50 (median)  | 0.50 | 1.25× |
/// | 0.70           | 0.79 | 1.40× |
/// | 0.95 (top impl)| 0.95 | 1.47× |
///
/// Ceiling at `1 + α` — with `α = 0.5` that's 1.5×, bounded enough to
/// keep PageRank a tiebreaker rather than a dominator: an irrelevant
/// top-PR file with dense score 0.6 gets `0.6 × 1.5 = 0.9` and still
/// loses to a relevant low-PR file scoring above 0.9.
///
/// This matches the two design constraints:
/// 1. A test (low percentile) should not be lifted above an impl
///    (high percentile) on similar dense scores. Sigmoid centered at
///    0.5 makes "below median" almost-no-boost.
/// 2. A heavily-imported file shouldn't dominate. The sigmoid plateau
///    above `percentile > 0.85` means a singularly-popular file gets
///    barely more boost than a moderately-popular one.
#[must_use]
pub fn pagerank_boost_factor(percentile: f32, alpha: f32) -> f32 {
    if percentile <= 0.0 || alpha <= 0.0 {
        return 1.0;
    }
    let z = (percentile.clamp(0.0, 1.0) - 0.5) / PAGERANK_SIGMOID_STEEPNESS;
    let sigmoid = 1.0 / (1.0 + (-z).exp());
    1.0 + alpha * sigmoid
}

/// Apply a multiplicative PageRank boost to search results.
///
/// For each result, looks up the chunk's PageRank percentile and applies
/// the sigmoid boost from [`pagerank_boost_factor`].
///
/// Results are re-sorted after boosting.
///
/// `pagerank_by_file` maps relative file paths to their **PageRank
/// percentile** in the corpus distribution — not the raw rank value.
/// Build it via [`pagerank_lookup`], which switched to percentile in
/// service of the sigmoid curve.
///
/// `alpha` controls the maximum boost (ceiling = `1 + alpha`). The
/// `alpha` field from [`RepoGraph`] is recommended (auto-tuned from
/// graph density).
pub fn boost_with_pagerank<S: std::hash::BuildHasher>(
    results: &mut [(usize, f32)],
    chunks: &[CodeChunk],
    pagerank_by_file: &HashMap<String, f32, S>,
    alpha: f32,
) {
    // Operates on `&mut [_]` (not `&mut Vec<_>`) so we can't delegate
    // to `crate::ranking::PageRankBoost::apply` directly (the trait
    // method takes `&mut Vec` to allow truncation layers). Replicate
    // the boost loop inline; both paths share `lookup_rank` +
    // `pagerank_boost_factor` so the curve stays consistent.
    for (idx, score) in results.iter_mut() {
        if let Some(chunk) = chunks.get(*idx) {
            let rank = lookup_rank(pagerank_by_file, &chunk.file_path, &chunk.name);
            *score *= pagerank_boost_factor(rank, alpha);
        }
    }
    results.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
}

/// `boost_with_pagerank` variant that operates on `SearchResult` directly,
/// for callers that don't have the raw `(usize, f32)` pair at hand.
///
/// Same boost math as [`boost_with_pagerank`]; re-sorts in place.
pub fn boost_with_pagerank_results<S: std::hash::BuildHasher>(
    results: &mut [crate::embed::SearchResult],
    pagerank_by_file: &HashMap<String, f32, S>,
    alpha: f32,
) {
    // SearchResult shape; inline math like `boost_with_pagerank`.
    for r in results.iter_mut() {
        let rank = lookup_rank(pagerank_by_file, &r.chunk.file_path, &r.chunk.name);
        r.similarity *= pagerank_boost_factor(rank, alpha);
    }
    results.sort_unstable_by(|a, b| b.similarity.total_cmp(&a.similarity));
}

/// Resolve a chunk's PageRank score from a path that may be rooted
/// differently than the graph keys.
///
/// Background: `RepoGraph` stores `FileNode.path` as `path.strip_prefix(root)`
/// where `root` is the **canonicalized** corpus root. Chunk
/// `file_path` is `path.display()` where `path` came from the walker —
/// which uses the caller-supplied root **as-is** (not canonicalized).
/// When the caller passes `tests/corpus/code/tokio`, chunk paths look
/// like `tests/corpus/code/tokio/tokio/src/.../foo.rs` while graph
/// keys look like `tokio/src/.../foo.rs`. Direct lookup never hits.
///
/// This function tries: definition-level exact (`"file::name"`),
/// file-level exact, then walks the chunk path one segment at a time
/// from the left and retries each suffix. First match wins.
///
/// (The proper fix is to normalize chunk paths at chunk-creation time
/// to be relative to the canonicalized corpus root; that's a larger
/// refactor planned alongside the `RankingLayer` work. Suffix matching
/// is the surgical patch that makes PageRank actually function.)
/// Re-exported under a longer name for use from the
/// [`crate::ranking`] module. Kept as a `pub(crate)` symbol so it
/// doesn't leak into the public surface; the canonical access point
/// is [`crate::ranking::PageRankBoost`].
pub(crate) fn lookup_rank_for_chunk<S: std::hash::BuildHasher>(
    pr: &HashMap<String, f32, S>,
    file_path: &str,
    name: &str,
) -> f32 {
    lookup_rank(pr, file_path, name)
}

fn lookup_rank<S: std::hash::BuildHasher>(
    pr: &HashMap<String, f32, S>,
    file_path: &str,
    name: &str,
) -> f32 {
    let def_key = format!("{file_path}::{name}");
    if let Some(&r) = pr.get(&def_key) {
        return r;
    }
    if let Some(&r) = pr.get(file_path) {
        return r;
    }
    // Slide a left-edge cursor through the path. For
    // `a/b/c/d/foo.rs` try `b/c/d/foo.rs`, then `c/d/foo.rs`, etc.
    // Path components are typically <= 8 levels, so this is cheap.
    let mut rest = file_path;
    while let Some(idx) = rest.find('/') {
        rest = &rest[idx + 1..];
        if rest.is_empty() {
            break;
        }
        let def_key = format!("{rest}::{name}");
        if let Some(&r) = pr.get(&def_key) {
            return r;
        }
        if let Some(&r) = pr.get(rest) {
            return r;
        }
    }
    0.0
}

/// Build a normalized PageRank lookup table from a [`RepoGraph`].
///
/// Returns a map from `"file_path::def_name"` to definition-level PageRank
/// normalized to `[0, 1]`. Also inserts file-level entries (`"file_path"`)
/// as aggregated fallback for chunks that don't match a specific definition.
#[must_use]
pub fn pagerank_lookup(graph: &crate::repo_map::RepoGraph) -> HashMap<String, f32> {
    // Switched from `rank / max_rank` (proportional) to percentile in
    // the corpus distribution. Rationale: a top-K result set typically
    // contains files whose raw ranks are all in a tiny band near zero
    // (Tokio: max in top-10 was 0.028 out of 1.0). Proportional
    // normalization gave uniformly tiny boosts. Percentile separates
    // "bottom decile (tests, leaves)" from "top half (impls, hubs)"
    // crisply, and pairs with the sigmoid in `pagerank_boost_factor`
    // to put the rank-transition where the action is.
    //
    // Definition-level and file-level percentiles use independent
    // distributions: `def_ranks` and `base_ranks`. A file that has no
    // defs still gets a file-level percentile from `base_ranks`.
    let def_pct = make_percentile_fn(&graph.def_ranks);
    let base_pct = make_percentile_fn(&graph.base_ranks);
    let mut map = HashMap::new();
    for (file_idx, file) in graph.files.iter().enumerate() {
        for (def_idx, def) in file.defs.iter().enumerate() {
            let flat = graph.def_offsets[file_idx] + def_idx;
            if let Some(&rank) = graph.def_ranks.get(flat) {
                let key = format!("{}::{}", file.path, def.name);
                map.insert(key, def_pct(rank));
            }
        }
        if file_idx < graph.base_ranks.len() {
            map.insert(file.path.clone(), base_pct(graph.base_ranks[file_idx]));
        }
    }
    map
}

/// Build a `value → percentile` function from a slice of rank values.
///
/// Sorts a copy once at build time, then each lookup is a binary search
/// over the sorted slice. Returns the empirical CDF: the fraction of
/// values strictly less than the queried value. Handles empty input
/// and `NaN` defensively.
fn make_percentile_fn(values: &[f32]) -> impl Fn(f32) -> f32 + '_ {
    let mut sorted: Vec<f32> = values.iter().copied().filter(|v| v.is_finite()).collect();
    sorted.sort_unstable_by(f32::total_cmp);
    move |value: f32| {
        if sorted.is_empty() {
            return 0.0;
        }
        // partition_point returns the count of elements strictly less
        // than `value` (because the predicate is `<`).
        let count_below = sorted.partition_point(|&v| v < value);
        #[expect(
            clippy::cast_precision_loss,
            reason = "rank counts well below f32 precision threshold"
        )]
        let pct = count_below as f32 / sorted.len() as f32;
        pct
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rrf_union_semantics() {
        // sem: [0, 1, 2], bm25: [3, 0, 4]
        // Chunk 0 appears in both lists → highest RRF score.
        // Chunks 1, 2, 3, 4 appear in exactly one list → all five appear.
        let sem = vec![(0, 0.9), (1, 0.8), (2, 0.7)];
        let bm25 = vec![(3, 10.0), (0, 8.0), (4, 6.0)];

        let fused = rrf_fuse(&sem, &bm25, 60.0);

        let indices: Vec<usize> = fused.iter().map(|&(i, _)| i).collect();

        // All 5 unique chunks must appear
        for expected in [0, 1, 2, 3, 4] {
            assert!(
                indices.contains(&expected),
                "chunk {expected} missing from fused results"
            );
        }
        assert_eq!(fused.len(), 5);

        // Chunk 0 must rank first (double-list bonus)
        assert_eq!(indices[0], 0, "chunk 0 should rank first");
    }

    #[test]
    fn rrf_single_list() {
        // Only semantic results; BM25 is empty.
        let sem = vec![(0, 0.9), (1, 0.8)];
        let bm25: Vec<(usize, f32)> = vec![];

        let fused = rrf_fuse(&sem, &bm25, 60.0);

        assert_eq!(fused.len(), 2);
        // Chunk 0 ranked first in sem list → higher RRF score than chunk 1
        assert_eq!(fused[0].0, 0);
        assert_eq!(fused[1].0, 1);
        assert!(fused[0].1 > fused[1].1);
    }

    #[test]
    fn search_mode_roundtrip() {
        assert_eq!("hybrid".parse::<SearchMode>().unwrap(), SearchMode::Hybrid);
        assert_eq!(
            "semantic".parse::<SearchMode>().unwrap(),
            SearchMode::Semantic
        );
        assert_eq!(
            "keyword".parse::<SearchMode>().unwrap(),
            SearchMode::Keyword
        );

        let err = "invalid".parse::<SearchMode>();
        assert!(err.is_err(), "expected parse error for 'invalid'");
        let msg = err.unwrap_err().to_string();
        assert!(
            msg.contains("invalid"),
            "error message should echo the bad input"
        );
    }

    #[test]
    fn search_mode_display() {
        assert_eq!(SearchMode::Hybrid.to_string(), "hybrid");
        assert_eq!(SearchMode::Semantic.to_string(), "semantic");
        assert_eq!(SearchMode::Keyword.to_string(), "keyword");
    }

    #[test]
    fn pagerank_boost_amplifies_relevant() {
        let chunks = vec![
            CodeChunk {
                file_path: "important.rs".into(),
                name: "a".into(),
                kind: "function".into(),
                start_line: 1,
                end_line: 10,
                content: String::new(),
                enriched_content: String::new(),
            },
            CodeChunk {
                file_path: "obscure.rs".into(),
                name: "b".into(),
                kind: "function".into(),
                start_line: 1,
                end_line: 10,
                content: String::new(),
                enriched_content: String::new(),
            },
        ];

        // Both start with same score; important.rs has high PageRank
        let mut results = vec![(0, 0.8_f32), (1, 0.8)];
        let mut pr = HashMap::new();
        pr.insert("important.rs".to_string(), 1.0); // max PageRank
        pr.insert("obscure.rs".to_string(), 0.1); // low PageRank

        boost_with_pagerank(&mut results, &chunks, &pr, 0.3);

        // important.rs should now rank higher
        assert_eq!(
            results[0].0, 0,
            "important.rs should rank first after boost"
        );
        assert!(results[0].1 > results[1].1);

        // Boost values reflect the sigmoid-on-percentile curve in
        // `pagerank_boost_factor` (alpha=0.3 here):
        // - percentile=1.0: sigmoid(3.33) ≈ 0.965, boost ≈ 1.29 → 1.032
        // - percentile=0.1: sigmoid(-2.67) ≈ 0.065, boost ≈ 1.02 → 0.816
        assert!(
            (results[0].1 - 1.032).abs() < 0.01,
            "rank=1.0 boost: expected ~1.032, got {}",
            results[0].1
        );
        assert!(
            (results[1].1 - 0.816).abs() < 0.01,
            "rank=0.1 boost: expected ~0.816, got {}",
            results[1].1
        );
    }

    #[test]
    fn pagerank_boost_zero_relevance_stays_zero() {
        let chunks = vec![CodeChunk {
            file_path: "important.rs".into(),
            name: "a".into(),
            kind: "function".into(),
            start_line: 1,
            end_line: 10,
            content: String::new(),
            enriched_content: String::new(),
        }];

        let mut results = vec![(0, 0.0_f32)];
        let mut pr = HashMap::new();
        pr.insert("important.rs".to_string(), 1.0);

        boost_with_pagerank(&mut results, &chunks, &pr, 0.3);

        // Zero score stays zero regardless of PageRank
        assert!(results[0].1.abs() < f32::EPSILON);
    }

    #[test]
    fn pagerank_boost_unknown_file_no_effect() {
        let chunks = vec![CodeChunk {
            file_path: "unknown.rs".into(),
            name: "a".into(),
            kind: "function".into(),
            start_line: 1,
            end_line: 10,
            content: String::new(),
            enriched_content: String::new(),
        }];

        let mut results = vec![(0, 0.5_f32)];
        let pr = HashMap::new(); // empty — no PageRank data

        boost_with_pagerank(&mut results, &chunks, &pr, 0.3);

        // No PageRank data → no boost
        assert!((results[0].1 - 0.5).abs() < f32::EPSILON);
    }
}