src/knowledge/fusion.rs
//! Reciprocal-Rank-Fusion (RRF) hybrid retrieval — Mimir SEARCH optimize cycle 1.
//!
//! Mimir already holds two independent, complementary retrieval signals over the
//! same code corpus, but until now they were never combined:
//!
//! * **BM25 full-text** — the persisted Tantivy index ([`crate::index::Index`]),
//! lexical/keyword relevance ("find the literal token `render_template`").
//! * **Dense cosine kNN** — the embedding store ([`crate::vector::store::search`]
//! over the `znippy_zoomies::vann` SIMD int8/f32 index), semantic relevance
//! ("find code that PARSES a template even if it never says `template`").
//!
//! Each alone misses what the other catches: BM25 whiffs on paraphrase; dense
//! kNN whiffs on rare exact identifiers/error-strings. **Reciprocal-Rank-Fusion**
//! merges the two *ranked* lists with no score calibration and no new dependency —
//! it fuses on RANK, so the wildly different score scales (BM25 tf-idf vs cosine
//! in `[-1,1]`) never need to be reconciled. The fused score of a key is
//!
//! ```text
//! rrf(key) = Σ_lists 1 / (k + rank_in_list(key)) // rank is 1-based
//! ```
//!
//! with `k` a smoothing constant (the literature default is 60). A key that ranks
//! *decently in both* lists beats a key that ranks #1 in only one — the consensus
//! property that lifts retrieval recall/precision at ~zero latency cost (one extra
//! rank walk over two lists we already computed).
//!
//! This module is **pure and engine-agnostic**: it fuses `Vec<String>` key lists,
//! so it is unit-testable without spinning a Tantivy index or an ONNX embedder.
//! The retrieval call sites (`knowledge_context` / the MCP search tool) fetch the
//! two lists from the real engines and hand them here.
/// The literature-default RRF smoothing constant. Larger `k` flattens the
/// contribution of top ranks (more democratic across the list); 60 is the value
/// from Cormack et al. 2009 and what every mainstream hybrid retriever ships.
pub const DEFAULT_RRF_K: f32 = 60.0;
/// A ranked retrieval result list tagged with an optional weight. `keys` are
/// best-first (rank 1 = index 0). `weight` scales this list's whole RRF
/// contribution (default `1.0`) so a caller can trust, say, dense kNN a little
/// more than lexical BM25 without touching the raw scores.
#[derive(Debug, Clone)]
pub struct RankedList<'a> {
pub keys: &'a [String],
pub weight: f32,
}
impl<'a> RankedList<'a> {
/// A list at the default weight (`1.0`).
pub fn new(keys: &'a [String]) -> Self {
RankedList { keys, weight: 1.0 }
}
/// A list at an explicit weight.
pub fn weighted(keys: &'a [String], weight: f32) -> Self {
RankedList { keys, weight }
}
}
/// Fuse many ranked key-lists into one ranking by Reciprocal-Rank-Fusion.
///
/// Returns `(key, fused_score)` best-first (descending score). A key present in
/// several lists accumulates `weight / (rrf_k + rank)` from each. Ties break on
/// the key string so the output is deterministic. `rrf_k` should be positive
/// (use [`DEFAULT_RRF_K`]); a non-positive value is clamped to a tiny epsilon so
/// the reciprocal never divides by zero.
///
/// The first appearance of a key in a list wins its rank (a well-formed ranked
/// list has no duplicates, but we are defensive: a repeated key does not get a
/// second, worse-ranked contribution from the same list).
pub fn reciprocal_rank_fusion(lists: &[RankedList<'_>], rrf_k: f32) -> Vec<(String, f32)> {
use std::collections::HashMap;
let k = if rrf_k > 0.0 { rrf_k } else { f32::EPSILON };
let mut scores: HashMap<&str, f32> = HashMap::new();
for list in lists {
let mut seen_in_list: HashMap<&str, ()> = HashMap::new();
for (i, key) in list.keys.iter().enumerate() {
// First occurrence in THIS list claims the rank; later dupes skipped.
if seen_in_list.insert(key.as_str(), ()).is_some() {
continue;
}
let rank = (i + 1) as f32; // 1-based
*scores.entry(key.as_str()).or_insert(0.0) += list.weight / (k + rank);
}
}
let mut out: Vec<(String, f32)> = scores
.into_iter()
.map(|(k, v)| (k.to_string(), v))
.collect();
// Best-first; deterministic tie-break on the key.
out.sort_by(|a, b| {
b.1.partial_cmp(&a.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.0.cmp(&b.0))
});
out
}
/// Convenience over [`reciprocal_rank_fusion`] for the common two-engine hybrid:
/// fuse a lexical (BM25) key list and a dense (cosine kNN) key list at the given
/// weights, returning the top-`limit` fused keys (best-first). The keys must be
/// in the SAME namespace in both lists (e.g. both repo-relative file paths) for
/// the consensus to line up. `limit == 0` returns the full fused ranking.
pub fn fuse_bm25_dense(
bm25_keys: &[String],
dense_keys: &[String],
bm25_weight: f32,
dense_weight: f32,
rrf_k: f32,
limit: usize,
) -> Vec<String> {
let lists = [
RankedList::weighted(bm25_keys, bm25_weight),
RankedList::weighted(dense_keys, dense_weight),
];
let mut fused: Vec<String> = reciprocal_rank_fusion(&lists, rrf_k)
.into_iter()
.map(|(k, _)| k)
.collect();
if limit > 0 && fused.len() > limit {
fused.truncate(limit);
}
fused
}
/// Recall@k of a `retrieved` ranking against a `relevant` set: the fraction of the
/// relevant keys that appear in the first `k` retrieved. The quality metric the
/// hybrid-vs-single-engine bench reports (RRF must not LOWER recall — the
/// red-when-broken guard on the fusion). `relevant` empty ⇒ `1.0` (vacuously
/// perfect); `k == 0` ⇒ `0.0`.
pub fn recall_at_k(retrieved: &[String], relevant: &[String], k: usize) -> f32 {
if relevant.is_empty() {
return 1.0;
}
if k == 0 {
return 0.0;
}
use std::collections::HashSet;
let top: HashSet<&str> = retrieved.iter().take(k).map(|s| s.as_str()).collect();
let hit = relevant.iter().filter(|r| top.contains(r.as_str())).count();
hit as f32 / relevant.len() as f32
}
#[cfg(test)]
mod tests {
use super::*;
fn keys(v: &[&str]) -> Vec<String> {
v.iter().map(|s| s.to_string()).collect()
}
/// RED-WHEN-BROKEN #1 — the CONSENSUS property: a key that ranks #2 in BOTH
/// lists must beat a key that ranks #1 in only ONE. This is the entire reason
/// RRF lifts hybrid recall. Break the reciprocal accumulation (e.g. `max`
/// instead of `sum`, or drop the second list) and `consensus` no longer wins.
#[test]
fn consensus_beats_single_list_number_one() {
let bm25 = keys(&["A", "consensus"]); // A@1, consensus@2
let dense = keys(&["X", "consensus"]); // X@1, consensus@2
let fused = fuse_bm25_dense(&bm25, &dense, 1.0, 1.0, DEFAULT_RRF_K, 0);
assert_eq!(
fused[0], "consensus",
"a decent-in-both key tops single-list #1s: {fused:?}"
);
// A and X each appear once at rank 1; consensus appears twice at rank 2.
// 1/(60+1) = 0.01639 (single) < 2/(60+2) = 0.03226 (consensus).
let pos = |k: &str| fused.iter().position(|x| x == k).unwrap();
assert!(pos("consensus") < pos("A") && pos("consensus") < pos("X"));
}
/// RED-WHEN-BROKEN #2 — a key ranked #1 in BOTH lists is the fused top; and
/// the fusion is order-stable / deterministic.
#[test]
fn agreed_top_is_first_and_deterministic() {
let bm25 = keys(&["top", "b", "c"]);
let dense = keys(&["top", "c", "b"]);
let f1 = fuse_bm25_dense(&bm25, &dense, 1.0, 1.0, DEFAULT_RRF_K, 0);
let f2 = fuse_bm25_dense(&bm25, &dense, 1.0, 1.0, DEFAULT_RRF_K, 0);
assert_eq!(f1[0], "top", "agreed #1 wins");
assert_eq!(f1, f2, "deterministic");
}
/// RED-WHEN-BROKEN #3 — the recall guard the bench asserts: fusing two
/// PARTIAL engines recovers relevant items that EITHER engine alone misses in
/// its top-k, so hybrid recall >= each single engine. Break fusion and this
/// inequality flips.
#[test]
fn hybrid_recall_at_least_each_single_engine() {
// Relevant set r1..r3. BM25 finds r1 high but buries r3; dense finds r3
// high but buries r1. Neither's top-2 covers all; the fusion's does better.
let relevant = keys(&["r1", "r2", "r3"]);
let bm25 = keys(&["r1", "r2", "n1", "n2", "r3"]);
let dense = keys(&["r3", "r2", "n3", "n4", "r1"]);
let k = 3;
let r_bm25 = recall_at_k(&bm25, &relevant, k);
let r_dense = recall_at_k(&dense, &relevant, k);
let fused = fuse_bm25_dense(&bm25, &dense, 1.0, 1.0, DEFAULT_RRF_K, 0);
let r_fused = recall_at_k(&fused, &relevant, k);
assert!(
r_fused >= r_bm25 && r_fused >= r_dense,
"hybrid recall {r_fused} must dominate bm25 {r_bm25} / dense {r_dense}; fused={fused:?}"
);
assert!(
r_fused > r_bm25 || r_fused > r_dense,
"and strictly beats at least one"
);
}
#[test]
fn weights_and_limit_apply() {
let bm25 = keys(&["a", "b"]);
let dense = keys(&["b", "a"]);
// Heavily up-weight dense: its rank-1 ("b") should win.
let fused = fuse_bm25_dense(&bm25, &dense, 0.1, 10.0, DEFAULT_RRF_K, 1);
assert_eq!(fused.len(), 1, "limit truncates");
assert_eq!(fused[0], "b", "dense weight dominates");
}
#[test]
fn recall_edges() {
assert_eq!(recall_at_k(&[], &[], 5), 1.0, "no relevant ⇒ vacuous 1.0");
assert_eq!(recall_at_k(&keys(&["a"]), &keys(&["a"]), 0), 0.0, "k=0 ⇒ 0");
assert_eq!(recall_at_k(&keys(&["a", "b"]), &keys(&["a", "b"]), 2), 1.0);
}
}