clawhdf5/crates/clawhdf5-agent/src/hybrid.rs

//! Hybrid search combining vector similarity and BM25 keyword scores.
//!
//! Normalizes both score sets to [0, 1] and computes a weighted merge.

use std::collections::HashMap;

use crate::bm25::BM25Index;
use crate::vector_search;

/// Perform hybrid search combining cosine vector similarity and BM25 keyword search.
///
/// Both score distributions are independently normalized to [0, 1] before
/// being combined with the specified weights. Uses pre-computed norms when
/// available for faster vector search.
///
/// # Arguments
///
/// * `query_embedding` - The query vector for cosine similarity.
/// * `query_text` - The query text for BM25 keyword search.
/// * `vectors` - All stored embedding vectors.
/// * `_chunks` - All stored text chunks (parallel to `vectors`).
/// * `tombstones` - Tombstone flags (non-zero = deleted).
/// * `bm25_index` - Pre-built BM25 index.
/// * `vector_weight` - Weight for vector similarity scores (default 0.7).
/// * `keyword_weight` - Weight for keyword search scores (default 0.3).
/// * `k` - Number of top results to return.
#[allow(clippy::too_many_arguments)]
pub fn hybrid_search(
    query_embedding: &[f32],
    query_text: &str,
    vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
    chunks: &[String],
    tombstones: &[u8],
    bm25_index: &BM25Index,
    vector_weight: f32,
    keyword_weight: f32,
    k: usize,
) -> Vec<(usize, f32)> {
    hybrid_search_fused(
        query_embedding,
        query_text,
        vectors,
        chunks,
        tombstones,
        bm25_index,
        Fusion::Weighted {
            vector: vector_weight,
            keyword: keyword_weight,
        },
        k,
    )
}

/// [`hybrid_search`] with the fusion method chosen explicitly.
#[allow(clippy::too_many_arguments)]
pub fn hybrid_search_fused(
    query_embedding: &[f32],
    query_text: &str,
    vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
    _chunks: &[String],
    tombstones: &[u8],
    bm25_index: &BM25Index,
    fusion: Fusion,
    k: usize,
) -> Vec<(usize, f32)> {
    // Get raw scores from both systems. Request all results so normalization
    // covers the full distribution.
    // Use parallel search when rayon feature is enabled and vector count > 10K.
    let vec_scores = {
        #[cfg(feature = "parallel")]
        {
            if vectors.count() > 10_000 {
                vector_search::parallel_cosine_batch(
                    query_embedding,
                    vectors,
                    tombstones,
                    vectors.count(),
                )
            } else {
                vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
            }
        }
        #[cfg(not(feature = "parallel"))]
        {
            vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
        }
    };
    let kw_scores = bm25_index.scores(query_text);

    fuse(vec_scores, kw_scores, fusion, k)
}

/// Merge pre-computed vector-similarity and keyword scores into a single ranking.
///
/// Both score sets are independently min-max normalized to [0, 1] and combined
/// with the given weights. This is the shared core of [`hybrid_search`]; it is
/// also used by the optional HNSW path, which supplies vector scores from an
/// approximate-nearest-neighbour index instead of a full linear scan.
pub fn merge_vector_keyword(
    vec_scores: Vec<(usize, f32)>,
    kw_scores: Vec<(usize, f32)>,
    vector_weight: f32,
    keyword_weight: f32,
    k: usize,
) -> Vec<(usize, f32)> {
    fuse(
        vec_scores,
        kw_scores,
        Fusion::Weighted {
            vector: vector_weight,
            keyword: keyword_weight,
        },
        k,
    )
}

/// How the vector and keyword stages are combined into one ranking.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Fusion {
    /// Min-max normalise each stage over its own candidates, then take a
    /// weighted sum. Uses the *scores*, so a stage that separates its
    /// candidates sharply keeps that separation — and a stage whose candidates
    /// are all near-identical contributes little.
    Weighted {
        /// Weight on the vector stage.
        vector: f32,
        /// Weight on the keyword stage.
        keyword: f32,
    },
    /// Reciprocal rank fusion: each stage contributes `1 / (k + rank)`,
    /// ignoring score magnitudes entirely. Robust when the two stages'
    /// scores aren't comparable, at the cost of discarding confidence.
    Rrf {
        /// The rank-damping constant; 60 is the value from the original paper.
        k: f32,
    },
}

impl Default for Fusion {
    fn default() -> Self {
        DEFAULT_FUSION
    }
}

/// The fusion `hybrid_search` uses unless told otherwise.
///
/// The weights are not a guess: a sweep of every 0.1 step over the full
/// LongMemEval haystack (500 questions, real MiniLM embeddings) found the
/// long-standing 0.7/0.3 default *strictly dominated* — 0.4/0.6 is better at
/// Hit@1, Hit@5, Hit@10 and MRR, at both turn and session granularity. See
/// `BENCHMARKS.md`, "Weight sweep".
pub const DEFAULT_FUSION: Fusion = Fusion::Weighted {
    vector: 0.4,
    keyword: 0.6,
};

/// Combine one ranked candidate list from each stage into a single top-`k`.
///
/// Neither list need be sorted; both are consumed.
pub fn fuse(
    vec_scores: Vec<(usize, f32)>,
    kw_scores: Vec<(usize, f32)>,
    fusion: Fusion,
    k: usize,
) -> Vec<(usize, f32)> {
    let mut merged: HashMap<usize, f32> = HashMap::new();
    match fusion {
        Fusion::Weighted { vector, keyword } => {
            // Normalize each set to [0, 1].
            for (idx, score) in &normalize_scores(&vec_scores) {
                *merged.entry(*idx).or_insert(0.0) += vector * score;
            }
            for (idx, score) in &normalize_scores(&kw_scores) {
                *merged.entry(*idx).or_insert(0.0) += keyword * score;
            }
        }
        Fusion::Rrf { k: damping } => {
            for mut stage in [vec_scores, kw_scores] {
                // Rank 1 is the best score. Ties break by index so a stage's
                // contribution doesn't depend on the candidate order it
                // happened to be produced in.
                stage.sort_by(|a, b| {
                    b.1.partial_cmp(&a.1)
                        .unwrap_or(std::cmp::Ordering::Equal)
                        .then(a.0.cmp(&b.0))
                });
                for (rank, (idx, _)) in stage.iter().enumerate() {
                    *merged.entry(*idx).or_insert(0.0) += 1.0 / (damping + (rank + 1) as f32);
                }
            }
        }
    }

    let mut results: Vec<(usize, f32)> = merged.into_iter().collect();
    // Index tie-break: `merged` is a HashMap, so without it the ties that
    // survive differ from run to run.
    let by_score_then_id = |a: &(usize, f32), b: &(usize, f32)| {
        b.1.partial_cmp(&a.1)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(a.0.cmp(&b.0))
    };
    // Only the top k are wanted: partition them out, then order just those,
    // instead of sorting every candidate (the keyword side can be the corpus).
    if k == 0 {
        return Vec::new();
    }
    if results.len() > k {
        results.select_nth_unstable_by(k - 1, by_score_then_id);
        results.truncate(k);
    }
    results.sort_by(by_score_then_id);
    results
}

/// Normalize a set of scores to the [0, 1] range using min-max normalization.
///
/// If all scores are identical there is no spread to normalise: each entry
/// gets 1.0 when that score is positive (all equally the best match) and 0.0
/// otherwise (nothing matched).
fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> {
    if scores.is_empty() {
        return Vec::new();
    }

    let min = scores.iter().map(|(_, s)| *s).fold(f32::INFINITY, f32::min);
    let max = scores
        .iter()
        .map(|(_, s)| *s)
        .fold(f32::NEG_INFINITY, f32::max);

    let range = max - min;
    if range == 0.0 {
        // All candidates scored the same (including the single-candidate
        // case), so min-max has no spread to work with. They are all equally
        // the best match if that score is positive, and all non-matches
        // otherwise. This used to return 0.0 unconditionally, which erased a
        // lone perfect match from the fused score.
        let level = if max > 0.0 { 1.0 } else { 0.0 };
        return scores.iter().map(|(idx, _)| (*idx, level)).collect();
    }

    scores
        .iter()
        .map(|(idx, s)| (*idx, (s - min) / range))
        .collect()
}

/// Perform hybrid search using Reciprocal Rank Fusion (RRF).
///
/// RRF combines rankings from multiple retrieval systems without requiring
/// score normalization. Each result is scored as:
///
///   `score = Σ 1 / (k + rank_i)`
///
/// where `k = 60` (standard constant that dampens the impact of high ranks)
/// and `rank_i` is the 1-based rank of the document in retrieval system `i`.
///
/// Documents only present in one system still receive a partial score.
///
/// # Arguments
///
/// * `query_embedding` - The query vector for cosine similarity.
/// * `query_text` - The query text for BM25 keyword search.
/// * `vectors` - All stored embedding vectors.
/// * `_chunks` - All stored text chunks (parallel to `vectors`).
/// * `tombstones` - Tombstone flags (non-zero = deleted).
/// * `bm25_index` - Pre-built BM25 index.
/// * `k` - Number of top results to return.
#[allow(clippy::too_many_arguments)]
pub fn rrf_hybrid_search(
    query_embedding: &[f32],
    query_text: &str,
    vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized),
    _chunks: &[String],
    tombstones: &[u8],
    bm25_index: &BM25Index,
    k: usize,
) -> Vec<(usize, f32)> {
    const RRF_K: f32 = 60.0;

    // Retrieve all results from both systems sorted descending by score.
    let mut vec_scores = {
        #[cfg(feature = "parallel")]
        {
            if vectors.count() > 10_000 {
                vector_search::parallel_cosine_batch(
                    query_embedding,
                    vectors,
                    tombstones,
                    vectors.count(),
                )
            } else {
                vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
            }
        }
        #[cfg(not(feature = "parallel"))]
        {
            vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones)
        }
    };
    let mut kw_scores = bm25_index.search(query_text, vectors.count());

    // Sort both lists descending so rank 1 = best.
    vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    kw_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    // Accumulate RRF scores.
    let mut rrf_scores: HashMap<usize, f32> = HashMap::new();

    for (rank, (idx, _score)) in vec_scores.iter().enumerate() {
        let rrf = 1.0 / (RRF_K + (rank + 1) as f32);
        *rrf_scores.entry(*idx).or_insert(0.0) += rrf;
    }
    for (rank, (idx, _score)) in kw_scores.iter().enumerate() {
        let rrf = 1.0 / (RRF_K + (rank + 1) as f32);
        *rrf_scores.entry(*idx).or_insert(0.0) += rrf;
    }

    let mut results: Vec<(usize, f32)> = rrf_scores.into_iter().collect();
    results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    results.truncate(k);
    results
}

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

    fn make_test_data() -> (Vec<Vec<f32>>, Vec<String>, Vec<u8>, BM25Index) {
        // 4 documents with 3-dim embeddings
        let vectors = vec![
            vec![1.0, 0.0, 0.0], // doc 0: points in x direction
            vec![0.0, 1.0, 0.0], // doc 1: points in y direction
            vec![0.7, 0.7, 0.0], // doc 2: between x and y
            vec![0.0, 0.0, 1.0], // doc 3: points in z direction
        ];
        let chunks = vec![
            "rust programming language".to_string(),
            "python scripting language".to_string(),
            "rust and python comparison".to_string(),
            "javascript web development".to_string(),
        ];
        let tombstones = vec![0u8; 4];
        let bm25 = BM25Index::build(&chunks, &tombstones);
        (vectors, chunks, tombstones, bm25)
    }

    #[test]
    fn vector_only_search() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![1.0, 0.0, 0.0]; // points in x, should match doc 0

        let results = hybrid_search(
            &query_emb,
            "nonexistent_xyz",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            1.0, // vector only
            0.0, // no keyword
            4,
        );

        assert!(!results.is_empty());
        assert_eq!(
            results[0].0, 0,
            "doc 0 should be top match for x-direction query"
        );
    }

    #[test]
    fn keyword_only_search() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        // Use a zero vector so vector similarity contributes nothing meaningful
        let query_emb = vec![0.0, 0.0, 0.0];

        let results = hybrid_search(
            &query_emb,
            "rust programming",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            0.0, // no vector
            1.0, // keyword only
            4,
        );

        assert!(!results.is_empty());
        // Doc 0 ("rust programming language") should rank highest for "rust programming"
        assert_eq!(results[0].0, 0);
    }

    #[test]
    fn balanced_merge_ranking() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        // Query embedding close to doc 0, text query for "rust"
        let query_emb = vec![0.9, 0.1, 0.0];

        let results = hybrid_search(
            &query_emb,
            "rust",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            0.7,
            0.3,
            4,
        );

        assert!(!results.is_empty());
        // Doc 0 should rank high (good vector match + contains "rust")
        // Doc 2 should also appear (contains "rust" + decent vector match)
        let top_ids: Vec<usize> = results.iter().map(|(idx, _)| *idx).collect();
        assert!(top_ids.contains(&0), "doc 0 should appear in results");
        assert!(top_ids.contains(&2), "doc 2 should appear in results");
    }

    #[test]
    fn empty_results_when_no_data() {
        let vectors: Vec<Vec<f32>> = Vec::new();
        let chunks: Vec<String> = Vec::new();
        let tombstones: Vec<u8> = Vec::new();
        let bm25 = BM25Index::build(&chunks, &tombstones);

        let results = hybrid_search(
            &[],
            "anything",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            0.7,
            0.3,
            10,
        );

        assert!(results.is_empty());
    }

    #[test]
    fn normalize_scores_empty() {
        let result = normalize_scores(&[]);
        assert!(result.is_empty());
    }

    #[test]
    fn normalize_scores_single() {
        // A lone positive score is the best match there is, not a non-match.
        let result = normalize_scores(&[(0, 5.0)]);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].1, 1.0);
    }

    #[test]
    fn default_fusion_is_the_tuned_operating_point() {
        // A sweep over the full LongMemEval haystack found 0.7/0.3 strictly
        // dominated by 0.4/0.6 (BENCHMARKS.md). This guards the finding
        // against being quietly undone.
        assert_eq!(
            DEFAULT_FUSION,
            Fusion::Weighted {
                vector: 0.4,
                keyword: 0.6
            }
        );
    }

    #[test]
    fn rrf_rewards_agreement_between_the_stages_and_ignores_magnitudes() {
        // Doc 1 is second-best in both stages; doc 0 is best in one and absent
        // from the other. RRF prefers the doc both stages liked.
        let vec_scores = vec![(0, 100.0), (1, 0.9)];
        let kw_scores = vec![(2, 5.0), (1, 4.9)];
        let ranked = fuse(vec_scores, kw_scores, Fusion::Rrf { k: 60.0 }, 3);
        assert_eq!(ranked[0].0, 1, "{ranked:?}");

        // Scaling one stage's scores cannot change an RRF ranking, only the
        // order within that stage can.
        let a = fuse(
            vec![(0, 1.0), (1, 0.5)],
            vec![(1, 2.0), (0, 1.0)],
            Fusion::Rrf { k: 60.0 },
            2,
        );
        let b = fuse(
            vec![(0, 1e6), (1, -3.0)],
            vec![(1, 0.002), (0, 0.001)],
            Fusion::Rrf { k: 60.0 },
            2,
        );
        assert_eq!(
            a.iter().map(|r| r.0).collect::<Vec<_>>(),
            b.iter().map(|r| r.0).collect::<Vec<_>>()
        );
    }

    #[test]
    fn merge_top_k_matches_a_full_sort() {
        // Many ties (scores repeat) so the index tie-break is exercised.
        let vec_scores: Vec<(usize, f32)> = (0..300).map(|i| (i, ((i * 7) % 13) as f32)).collect();
        let kw_scores: Vec<(usize, f32)> = (100..500).map(|i| (i, ((i * 5) % 11) as f32)).collect();
        let everything =
            merge_vector_keyword(vec_scores.clone(), kw_scores.clone(), 0.7, 0.3, 10_000);
        assert_eq!(everything.len(), 500);
        assert!(
            everything
                .windows(2)
                .all(|w| { w[0].1 > w[1].1 || (w[0].1 == w[1].1 && w[0].0 < w[1].0) })
        );
        for k in [0, 1, 7, 50, 499, 500, 501] {
            let top = merge_vector_keyword(vec_scores.clone(), kw_scores.clone(), 0.7, 0.3, k);
            assert_eq!(top, everything[..k.min(500)], "k = {k}");
        }
    }

    #[test]
    fn normalize_scores_all_equal() {
        let matched = normalize_scores(&[(0, 0.4), (1, 0.4)]);
        assert!(matched.iter().all(|(_, s)| *s == 1.0));
        let unmatched = normalize_scores(&[(0, 0.0), (1, 0.0)]);
        assert!(unmatched.iter().all(|(_, s)| *s == 0.0));
    }

    #[test]
    fn normalize_scores_range() {
        let scores = vec![(0, 2.0), (1, 4.0), (2, 6.0)];
        let result = normalize_scores(&scores);

        assert_eq!(result.len(), 3);
        assert!((result[0].1 - 0.0).abs() < 1e-6); // min -> 0
        assert!((result[1].1 - 0.5).abs() < 1e-6); // mid -> 0.5
        assert!((result[2].1 - 1.0).abs() < 1e-6); // max -> 1
    }

    #[test]
    fn hybrid_respects_k_limit() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![0.5, 0.5, 0.0];

        let results = hybrid_search(
            &query_emb,
            "language",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            0.5,
            0.5,
            2,
        );

        assert!(results.len() <= 2);
    }

    // --- RRF tests ---

    #[test]
    fn rrf_vector_dominant_query() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![1.0, 0.0, 0.0]; // strong match on doc 0

        let results = rrf_hybrid_search(
            &query_emb,
            "nonexistent_xyz",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            4,
        );

        assert!(!results.is_empty());
        assert_eq!(
            results[0].0, 0,
            "doc 0 should top RRF for x-direction query"
        );
    }

    #[test]
    fn rrf_keyword_dominant_query() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![0.0, 0.0, 0.0];

        let results = rrf_hybrid_search(
            &query_emb,
            "rust programming",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            4,
        );

        assert!(!results.is_empty());
        assert_eq!(
            results[0].0, 0,
            "doc 0 should top RRF for rust programming query"
        );
    }

    #[test]
    fn rrf_respects_k_limit() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![0.5, 0.5, 0.0];

        let results = rrf_hybrid_search(
            &query_emb,
            "language",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            2,
        );

        assert!(results.len() <= 2);
    }

    #[test]
    fn rrf_scores_are_positive() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![0.5, 0.5, 0.0];

        let results =
            rrf_hybrid_search(&query_emb, "rust", &vectors, &chunks, &tombstones, &bm25, 4);

        for (_, score) in &results {
            assert!(*score > 0.0, "RRF scores must be positive");
        }
    }

    #[test]
    fn rrf_empty_data() {
        let vectors: Vec<Vec<f32>> = Vec::new();
        let chunks: Vec<String> = Vec::new();
        let tombstones: Vec<u8> = Vec::new();
        let bm25 = BM25Index::build(&chunks, &tombstones);

        let results = rrf_hybrid_search(&[], "anything", &vectors, &chunks, &tombstones, &bm25, 10);

        assert!(results.is_empty());
    }

    #[test]
    fn rrf_scores_sorted_descending() {
        let (vectors, chunks, tombstones, bm25) = make_test_data();
        let query_emb = vec![0.7, 0.3, 0.0];

        let results = rrf_hybrid_search(
            &query_emb,
            "rust language",
            &vectors,
            &chunks,
            &tombstones,
            &bm25,
            4,
        );

        for window in results.windows(2) {
            assert!(
                window[0].1 >= window[1].1,
                "RRF results must be sorted descending"
            );
        }
    }
}