348 lines
12 KiB
Rust
348 lines
12 KiB
Rust
|
//! Maximal Marginal Relevance (MMR) diversity re-ranking.
|
|||
|
|
//!
|
|||
|
|
//! Without MMR, if a user has multiple memory chunks about the same topic,
|
|||
|
|
//! the top results are nearly identical. MMR penalizes redundancy by
|
|||
|
|
//! greedily selecting results that balance relevance with diversity.
|
|||
|
|
//!
|
|||
|
|
//! **Formula:**
|
|||
|
|
//! ```text
|
|||
|
|
//! MMR(d) = λ × relevance(d) - (1-λ) × max_similarity(d, selected)
|
|||
|
|
//! ```
|
|||
|
|
//!
|
|||
|
|
//! Uses Jaccard similarity on tokenized snippets (no embeddings needed).
|
|||
|
|
//! O(n²) but n is tiny (typically 6–18 candidates after hybrid scoring).
|
|||
|
|
|
|||
|
|
use std::collections::HashSet;
|
|||
|
|
|
|||
|
|
use super::search::SearchResult;
|
|||
|
|
use xai_grok_config_types::MmrConfig;
|
|||
|
|
|
|||
|
|
/// Tokenize text into a set of alphanumeric words for Jaccard comparison.
|
|||
|
|
///
|
|||
|
|
/// Expects **pre-lowered** input — callers should lowercase snippets before
|
|||
|
|
/// calling this. Uses the same splitting strategy as `query_expansion`
|
|||
|
|
/// (split on non-alphanumeric except underscore) for consistency, but without
|
|||
|
|
/// stop word removal — we want full token overlap for similarity measurement.
|
|||
|
|
fn tokenize(text: &str) -> HashSet<&str> {
|
|||
|
|
text.split(|c: char| !c.is_alphanumeric() && c != '_')
|
|||
|
|
.filter(|w| !w.is_empty())
|
|||
|
|
.collect()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/// Jaccard similarity: |A ∩ B| / |A ∪ B|.
|
|||
|
|
fn jaccard_similarity(a: &HashSet<&str>, b: &HashSet<&str>) -> f64 {
|
|||
|
|
if a.is_empty() && b.is_empty() {
|
|||
|
|
return 1.0;
|
|||
|
|
}
|
|||
|
|
if a.is_empty() || b.is_empty() {
|
|||
|
|
return 0.0;
|
|||
|
|
}
|
|||
|
|
let intersection = a.intersection(b).count();
|
|||
|
|
let union = a.len() + b.len() - intersection;
|
|||
|
|
if union == 0 {
|
|||
|
|
0.0
|
|||
|
|
} else {
|
|||
|
|
intersection as f64 / union as f64
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/// Re-rank results using Maximal Marginal Relevance.
|
|||
|
|
///
|
|||
|
|
/// Reorders `results` in-place to balance relevance with diversity.
|
|||
|
|
/// No-op when `config.enabled` is false, `lambda` is 1.0, or there
|
|||
|
|
/// are fewer than 2 results.
|
|||
|
|
///
|
|||
|
|
/// `relevance` is the per-result unclamped ranking score, aligned
|
|||
|
|
/// index-for-index with `results` on entry. It is passed separately rather than
|
|||
|
|
/// read from the clamped `SearchResult.score`, which would saturate top chunks
|
|||
|
|
/// to 1.0 and lose the access-frequency boost tiebreak.
|
|||
|
|
pub fn mmr_rerank(results: &mut Vec<SearchResult>, relevance: &[f64], config: &MmrConfig) {
|
|||
|
|
if !config.enabled || results.len() <= 1 {
|
|||
|
|
return;
|
|||
|
|
}
|
|||
|
|
if config.lambda == 1.0 {
|
|||
|
|
return;
|
|||
|
|
}
|
|||
|
|
assert_eq!(
|
|||
|
|
relevance.len(),
|
|||
|
|
results.len(),
|
|||
|
|
"relevance must be aligned with results"
|
|||
|
|
);
|
|||
|
|
|
|||
|
|
// Lowercase snippets once, then tokenize. This ensures "Rust" and "rust"
|
|||
|
|
// are treated as the same token — casing varies across markdown sources.
|
|||
|
|
let lowered: Vec<String> = results.iter().map(|r| r.snippet.to_lowercase()).collect();
|
|||
|
|
let token_cache: Vec<HashSet<&str>> = lowered.iter().map(|s| tokenize(s)).collect();
|
|||
|
|
|
|||
|
|
let max_score = relevance.iter().copied().fold(f64::NEG_INFINITY, f64::max);
|
|||
|
|
let min_score = relevance.iter().copied().fold(f64::INFINITY, f64::min);
|
|||
|
|
let range = (max_score - min_score).max(f64::EPSILON);
|
|||
|
|
|
|||
|
|
let lambda = config.lambda;
|
|||
|
|
let mut selected: Vec<usize> = Vec::with_capacity(results.len());
|
|||
|
|
let mut remaining: Vec<usize> = (0..results.len()).collect();
|
|||
|
|
|
|||
|
|
while !remaining.is_empty() {
|
|||
|
|
let mut best_pos = 0;
|
|||
|
|
let mut best_mmr = f64::NEG_INFINITY;
|
|||
|
|
|
|||
|
|
for (pos, &candidate) in remaining.iter().enumerate() {
|
|||
|
|
let normalized = (relevance[candidate] - min_score) / range;
|
|||
|
|
|
|||
|
|
let max_sim = selected
|
|||
|
|
.iter()
|
|||
|
|
.map(|&sel| jaccard_similarity(&token_cache[candidate], &token_cache[sel]))
|
|||
|
|
.fold(0.0_f64, f64::max);
|
|||
|
|
|
|||
|
|
let mmr_score = lambda * normalized - (1.0 - lambda) * max_sim;
|
|||
|
|
|
|||
|
|
if mmr_score > best_mmr
|
|||
|
|
|| (mmr_score == best_mmr && relevance[candidate] > relevance[remaining[best_pos]])
|
|||
|
|
{
|
|||
|
|
best_mmr = mmr_score;
|
|||
|
|
best_pos = pos;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
selected.push(remaining.remove(best_pos));
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
let reordered: Vec<SearchResult> = selected
|
|||
|
|
.into_iter()
|
|||
|
|
.map(|i| std::mem::replace(&mut results[i], placeholder_result()))
|
|||
|
|
.collect();
|
|||
|
|
*results = reordered;
|
|||
|
|
// `results` is now reordered, so the caller's `relevance` slice is stale
|
|||
|
|
// and must not be read again.
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/// Placeholder to enable moving results out of the vec without Clone.
|
|||
|
|
fn placeholder_result() -> SearchResult {
|
|||
|
|
SearchResult {
|
|||
|
|
chunk_id: String::new(),
|
|||
|
|
path: String::new(),
|
|||
|
|
start_line: 0,
|
|||
|
|
end_line: 0,
|
|||
|
|
score: 0.0,
|
|||
|
|
snippet: String::new(),
|
|||
|
|
source: String::new(),
|
|||
|
|
created_at: 0,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[cfg(test)]
|
|||
|
|
mod tests {
|
|||
|
|
use super::*;
|
|||
|
|
|
|||
|
|
fn make_result(id: &str, snippet: &str, score: f64) -> SearchResult {
|
|||
|
|
SearchResult {
|
|||
|
|
chunk_id: id.to_string(),
|
|||
|
|
path: format!("{id}.md"),
|
|||
|
|
start_line: 0,
|
|||
|
|
end_line: 1,
|
|||
|
|
score,
|
|||
|
|
snippet: snippet.to_string(),
|
|||
|
|
source: "workspace".to_string(),
|
|||
|
|
created_at: 1_700_000_000,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
fn enabled_config(lambda: f64) -> MmrConfig {
|
|||
|
|
MmrConfig {
|
|||
|
|
enabled: true,
|
|||
|
|
lambda,
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/// Test helper: re-rank using each result's own `score` as its relevance
|
|||
|
|
/// (mirrors the pre-split behavior the existing assertions were written for).
|
|||
|
|
fn rerank(results: &mut Vec<SearchResult>, config: &MmrConfig) {
|
|||
|
|
let relevance: Vec<f64> = results.iter().map(|r| r.score).collect();
|
|||
|
|
mmr_rerank(results, &relevance, config);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_disabled_is_noop() {
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "rust async", 1.0),
|
|||
|
|
make_result("b", "rust async patterns", 0.9),
|
|||
|
|
];
|
|||
|
|
let original_order: Vec<String> = results.iter().map(|r| r.chunk_id.clone()).collect();
|
|||
|
|
rerank(&mut results, &MmrConfig::default());
|
|||
|
|
let after: Vec<String> = results.iter().map(|r| r.chunk_id.clone()).collect();
|
|||
|
|
assert_eq!(original_order, after);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_lambda_one_is_noop() {
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "rust async", 1.0),
|
|||
|
|
make_result("b", "python sync", 0.5),
|
|||
|
|
];
|
|||
|
|
rerank(&mut results, &enabled_config(1.0));
|
|||
|
|
assert_eq!(results[0].chunk_id, "a");
|
|||
|
|
assert_eq!(results[1].chunk_id, "b");
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_single_result_is_noop() {
|
|||
|
|
let mut results = vec![make_result("a", "rust async", 1.0)];
|
|||
|
|
rerank(&mut results, &enabled_config(0.7));
|
|||
|
|
assert_eq!(results.len(), 1);
|
|||
|
|
assert_eq!(results[0].chunk_id, "a");
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
/// Regression guard: MMR must rank on `relevance`, not the clamped
|
|||
|
|
/// `SearchResult.score`. Both results tie at `score == 1.0`; the
|
|||
|
|
/// higher-relevance result is placed SECOND so a buggy `.score` read would
|
|||
|
|
/// keep input order and land "low" first.
|
|||
|
|
#[test]
|
|||
|
|
fn test_mmr_ranks_on_relevance_not_clamped_score() {
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("low", "alpha topic one", 1.0),
|
|||
|
|
make_result("high", "beta subject two", 1.0),
|
|||
|
|
];
|
|||
|
|
let relevance = [1.0, 1.25];
|
|||
|
|
mmr_rerank(&mut results, &relevance, &enabled_config(0.7));
|
|||
|
|
|
|||
|
|
assert_eq!(
|
|||
|
|
results[0].chunk_id, "high",
|
|||
|
|
"MMR must order by unclamped relevance, not the clamped .score",
|
|||
|
|
);
|
|||
|
|
assert_eq!(results[1].chunk_id, "low");
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_diverse_results_promoted() {
|
|||
|
|
// Three results: two very similar (rust async), one different (python web)
|
|||
|
|
// With MMR, the diverse result should be promoted over the redundant one
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "rust async programming patterns", 1.0),
|
|||
|
|
make_result("b", "rust async programming tutorial", 0.95),
|
|||
|
|
make_result("c", "python web framework flask", 0.9),
|
|||
|
|
];
|
|||
|
|
rerank(&mut results, &enabled_config(0.5));
|
|||
|
|
|
|||
|
|
// First should still be "a" (highest relevance)
|
|||
|
|
assert_eq!(results[0].chunk_id, "a");
|
|||
|
|
// "c" (diverse) should be promoted above "b" (redundant with "a")
|
|||
|
|
assert_eq!(
|
|||
|
|
results[1].chunk_id, "c",
|
|||
|
|
"diverse result should be promoted over redundant one"
|
|||
|
|
);
|
|||
|
|
assert_eq!(results[2].chunk_id, "b");
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_identical_snippets_heavily_penalized() {
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "exact same content here", 1.0),
|
|||
|
|
make_result("b", "exact same content here", 0.99),
|
|||
|
|
make_result("c", "completely different topic", 0.5),
|
|||
|
|
];
|
|||
|
|
rerank(&mut results, &enabled_config(0.5));
|
|||
|
|
|
|||
|
|
assert_eq!(results[0].chunk_id, "a");
|
|||
|
|
// "c" should beat "b" because "b" is identical to "a"
|
|||
|
|
assert_eq!(
|
|||
|
|
results[1].chunk_id, "c",
|
|||
|
|
"different result should beat identical duplicate"
|
|||
|
|
);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_case_insensitive_similarity() {
|
|||
|
|
// "Rust Async" and "rust async" should be treated as identical
|
|||
|
|
// (both lowercased before tokenization). Without lowercasing,
|
|||
|
|
// these would only have 0.5 Jaccard similarity.
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "Rust Async Programming", 1.0),
|
|||
|
|
make_result("b", "rust async programming", 0.95),
|
|||
|
|
make_result("c", "Python Web Framework", 0.9),
|
|||
|
|
];
|
|||
|
|
rerank(&mut results, &enabled_config(0.5));
|
|||
|
|
|
|||
|
|
assert_eq!(results[0].chunk_id, "a");
|
|||
|
|
// "c" (diverse) should beat "b" (same content, different casing)
|
|||
|
|
assert_eq!(
|
|||
|
|
results[1].chunk_id, "c",
|
|||
|
|
"case-only difference should be detected as redundant"
|
|||
|
|
);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_preserves_result_count() {
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "one", 1.0),
|
|||
|
|
make_result("b", "two", 0.9),
|
|||
|
|
make_result("c", "three", 0.8),
|
|||
|
|
make_result("d", "four", 0.7),
|
|||
|
|
];
|
|||
|
|
rerank(&mut results, &enabled_config(0.7));
|
|||
|
|
assert_eq!(results.len(), 4);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_scores_and_snippets_preserved() {
|
|||
|
|
let mut results = vec![
|
|||
|
|
make_result("a", "rust programming", 1.0),
|
|||
|
|
make_result("b", "python scripting", 0.5),
|
|||
|
|
];
|
|||
|
|
rerank(&mut results, &enabled_config(0.7));
|
|||
|
|
// All fields should be intact after re-ranking
|
|||
|
|
for r in &results {
|
|||
|
|
assert!(!r.chunk_id.is_empty());
|
|||
|
|
assert!(!r.snippet.is_empty());
|
|||
|
|
assert!(r.score > 0.0);
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// -----------------------------------------------------------------------
|
|||
|
|
// Jaccard similarity unit tests
|
|||
|
|
// -----------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_jaccard_identical() {
|
|||
|
|
let a: HashSet<&str> = ["rust", "async"].into();
|
|||
|
|
let b: HashSet<&str> = ["rust", "async"].into();
|
|||
|
|
assert!((jaccard_similarity(&a, &b) - 1.0).abs() < f64::EPSILON);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_jaccard_disjoint() {
|
|||
|
|
let a: HashSet<&str> = ["rust", "async"].into();
|
|||
|
|
let b: HashSet<&str> = ["python", "web"].into();
|
|||
|
|
assert!((jaccard_similarity(&a, &b)).abs() < f64::EPSILON);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_jaccard_partial_overlap() {
|
|||
|
|
let a: HashSet<&str> = ["rust", "async", "programming"].into();
|
|||
|
|
let b: HashSet<&str> = ["rust", "web", "programming"].into();
|
|||
|
|
// intersection = {rust, programming} = 2, union = {rust, async, programming, web} = 4
|
|||
|
|
assert!((jaccard_similarity(&a, &b) - 0.5).abs() < f64::EPSILON);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_jaccard_both_empty() {
|
|||
|
|
let a: HashSet<&str> = HashSet::new();
|
|||
|
|
let b: HashSet<&str> = HashSet::new();
|
|||
|
|
assert!((jaccard_similarity(&a, &b) - 1.0).abs() < f64::EPSILON);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_jaccard_one_empty() {
|
|||
|
|
let a: HashSet<&str> = ["rust"].into();
|
|||
|
|
let b: HashSet<&str> = HashSet::new();
|
|||
|
|
assert!((jaccard_similarity(&a, &b)).abs() < f64::EPSILON);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
#[test]
|
|||
|
|
fn test_tokenize_splits_on_punctuation() {
|
|||
|
|
let tokens = tokenize("hello, world! rust_code");
|
|||
|
|
assert!(tokens.contains("hello"));
|
|||
|
|
assert!(tokens.contains("world"));
|
|||
|
|
assert!(tokens.contains("rust_code"));
|
|||
|
|
assert!(!tokens.contains(","));
|
|||
|
|
}
|
|||
|
|
}
|