use std::collections::HashMap;
use std::fmt;
#[derive(Debug, Clone, PartialEq)]
pub enum RerankerError {
NoCandidates,
IncompatibleDimensions {
expected: usize,
got: usize,
},
InvalidWeight(f64),
ConfigurationError(String),
}
impl fmt::Display for RerankerError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
RerankerError::NoCandidates => {
write!(f, "no candidates provided for reranking")
}
RerankerError::IncompatibleDimensions { expected, got } => {
write!(
f,
"embedding dimension mismatch: expected {expected}, got {got}"
)
}
RerankerError::InvalidWeight(w) => {
write!(f, "invalid weight value: {w}")
}
RerankerError::ConfigurationError(msg) => {
write!(f, "configuration error: {msg}")
}
}
}
}
impl std::error::Error for RerankerError {}
#[derive(Debug, Clone)]
pub struct ModalityScore {
pub modality: String,
pub raw_score: f64,
pub normalized_score: f64,
pub weight: f64,
}
impl ModalityScore {
pub fn new(modality: impl Into<String>, raw_score: f64, weight: f64) -> Self {
Self {
modality: modality.into(),
raw_score,
normalized_score: raw_score,
weight,
}
}
}
#[derive(Debug, Clone)]
pub struct RerankerCandidate {
pub id: String,
pub text_snippet: Option<String>,
pub embedding: Option<Vec<f64>>,
pub modality_scores: Vec<ModalityScore>,
pub final_score: f64,
pub rank: usize,
}
impl RerankerCandidate {
pub fn new(
id: impl Into<String>,
text_snippet: Option<&str>,
embedding: Option<Vec<f64>>,
) -> Self {
Self {
id: id.into(),
text_snippet: text_snippet.map(str::to_owned),
embedding,
modality_scores: Vec::new(),
final_score: 0.0,
rank: 0,
}
}
}
#[derive(Debug, Clone)]
pub struct TextFeatures {
pub term_frequency: Vec<(String, f64)>,
pub bm25_score: f64,
pub exact_match_bonus: f64,
pub length_penalty: f64,
}
#[derive(Debug, Clone)]
pub struct VectorFeatures {
pub cosine_similarity: f64,
pub dot_product: f64,
pub l2_distance: f64,
pub euclidean_normalized: f64,
}
#[derive(Debug, Clone)]
pub enum CmrFusionStrategy {
LinearCombination(Vec<(String, f64)>),
ReciprocalRankFusion(f64),
Borda,
MaxScore,
LearnedWeights(Vec<f64>),
}
impl Default for CmrFusionStrategy {
fn default() -> Self {
CmrFusionStrategy::LinearCombination(vec![
("text".to_string(), 0.4),
("vector".to_string(), 0.6),
])
}
}
#[derive(Debug, Clone)]
pub struct RerankerConfig {
pub fusion_strategy: CmrFusionStrategy,
pub text_weight: f64,
pub vector_weight: f64,
pub bm25_k1: f64,
pub bm25_b: f64,
pub normalize_scores: bool,
pub min_score_threshold: f64,
pub top_k: usize,
}
impl Default for RerankerConfig {
fn default() -> Self {
Self {
fusion_strategy: CmrFusionStrategy::default(),
text_weight: 0.4,
vector_weight: 0.6,
bm25_k1: 1.5,
bm25_b: 0.75,
normalize_scores: true,
min_score_threshold: 0.0,
top_k: 100,
}
}
}
impl RerankerConfig {
fn validate(&self) -> Result<(), RerankerError> {
if self.text_weight < 0.0 || self.text_weight.is_nan() {
return Err(RerankerError::InvalidWeight(self.text_weight));
}
if self.vector_weight < 0.0 || self.vector_weight.is_nan() {
return Err(RerankerError::InvalidWeight(self.vector_weight));
}
if self.bm25_k1 < 0.0 || self.bm25_k1.is_nan() {
return Err(RerankerError::ConfigurationError(
"bm25_k1 must be non-negative".to_string(),
));
}
if !(0.0..=1.0).contains(&self.bm25_b) {
return Err(RerankerError::ConfigurationError(
"bm25_b must be in [0, 1]".to_string(),
));
}
if let CmrFusionStrategy::LinearCombination(ref pairs) = self.fusion_strategy {
for (_, w) in pairs {
if *w < 0.0 || w.is_nan() {
return Err(RerankerError::InvalidWeight(*w));
}
}
}
if let CmrFusionStrategy::ReciprocalRankFusion(k) = self.fusion_strategy {
if k <= 0.0 || k.is_nan() {
return Err(RerankerError::ConfigurationError(
"RRF k must be positive".to_string(),
));
}
}
Ok(())
}
}
#[derive(Debug, Clone, Default)]
pub struct RerankerStats {
pub candidates_reranked: u64,
pub avg_rank_displacement: f64,
pub modalities_used: Vec<String>,
pub fusion_calls: u64,
}
fn tokenize(text: &str) -> Vec<String> {
text.split_whitespace()
.map(|w| {
w.to_lowercase()
.trim_matches(|c: char| !c.is_alphabetic())
.to_string()
})
.filter(|w| !w.is_empty())
.collect()
}
#[allow(dead_code)]
fn xorshift64(state: &mut u64) -> u64 {
let mut x = *state;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
*state = x;
x
}
#[allow(dead_code)]
fn xorshift_f64(state: &mut u64) -> f64 {
(xorshift64(state) >> 11) as f64 / (1u64 << 53) as f64
}
pub struct CrossModalReranker {
config: RerankerConfig,
stats: RerankerStats,
}
impl CrossModalReranker {
pub fn new(config: RerankerConfig) -> Self {
Self {
config,
stats: RerankerStats::default(),
}
}
pub fn update_config(&mut self, config: RerankerConfig) {
self.config = config;
}
pub fn stats(&self) -> RerankerStats {
self.stats.clone()
}
pub fn compute_text_features(&self, query: &str, text: &str, avg_doc_len: f64) -> TextFeatures {
let avg_doc_len = if avg_doc_len > 0.0 { avg_doc_len } else { 1.0 };
let query_tokens = tokenize(query);
let doc_tokens = tokenize(text);
let doc_len = doc_tokens.len() as f64;
let mut tf_map: HashMap<String, f64> = HashMap::new();
for tok in &doc_tokens {
*tf_map.entry(tok.clone()).or_insert(0.0) += 1.0;
}
let k1 = self.config.bm25_k1;
let b = self.config.bm25_b;
let n = 1.0_f64;
let mut term_contributions: Vec<(String, f64)> = Vec::new();
let mut bm25_total = 0.0_f64;
for qt in &query_tokens {
let freq = tf_map.get(qt).copied().unwrap_or(0.0);
let df = if freq > 0.0 { 1.0 } else { 0.0 };
let idf = ((n - df + 0.5) / (df + 0.5) + 1.0).ln();
let tf_norm = (freq * (k1 + 1.0)) / (freq + k1 * (1.0 - b + b * doc_len / avg_doc_len));
let contribution = idf * tf_norm;
bm25_total += contribution;
term_contributions.push((qt.clone(), contribution));
}
let exact_match_bonus =
if !query.is_empty() && text.to_lowercase().contains(&query.to_lowercase()) {
0.5
} else {
0.0
};
let length_penalty = 1.0 - 0.1 * ((doc_len / avg_doc_len) - 1.0).max(0.0);
TextFeatures {
term_frequency: term_contributions,
bm25_score: bm25_total,
exact_match_bonus,
length_penalty,
}
}
pub fn compute_vector_features(
query: &[f64],
candidate: &[f64],
) -> Result<VectorFeatures, RerankerError> {
if query.len() != candidate.len() {
return Err(RerankerError::IncompatibleDimensions {
expected: query.len(),
got: candidate.len(),
});
}
let mut dot = 0.0_f64;
let mut norm_q = 0.0_f64;
let mut norm_c = 0.0_f64;
let mut sq_diff = 0.0_f64;
for (q, c) in query.iter().zip(candidate.iter()) {
dot += q * c;
norm_q += q * q;
norm_c += c * c;
let d = q - c;
sq_diff += d * d;
}
let norm_q = norm_q.sqrt();
let norm_c = norm_c.sqrt();
let denom = norm_q * norm_c;
let cosine_similarity = if denom > 0.0 { dot / denom } else { 0.0 };
let l2_distance = sq_diff.sqrt();
let euclidean_normalized = 1.0 / (1.0 + l2_distance);
Ok(VectorFeatures {
cosine_similarity,
dot_product: dot,
l2_distance,
euclidean_normalized,
})
}
pub fn reciprocal_rank_fusion(rank_lists: Vec<Vec<String>>, k: f64) -> Vec<(String, f64)> {
let mut scores: HashMap<String, f64> = HashMap::new();
for list in &rank_lists {
for (rank_zero_based, id) in list.iter().enumerate() {
let rank = (rank_zero_based + 1) as f64;
*scores.entry(id.clone()).or_insert(0.0) += 1.0 / (k + rank);
}
}
let mut result: Vec<(String, f64)> = scores.into_iter().collect();
result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
result
}
pub fn rerank(
&mut self,
mut candidates: Vec<RerankerCandidate>,
query_text: Option<&str>,
query_embedding: Option<&[f64]>,
) -> Result<Vec<RerankerCandidate>, RerankerError> {
if candidates.is_empty() {
return Err(RerankerError::NoCandidates);
}
self.config.validate()?;
if let Some(qe) = query_embedding {
for c in &candidates {
if let Some(ce) = &c.embedding {
if ce.len() != qe.len() {
return Err(RerankerError::IncompatibleDimensions {
expected: qe.len(),
got: ce.len(),
});
}
}
}
}
let avg_doc_len = {
let texts: Vec<usize> = candidates
.iter()
.filter_map(|c| c.text_snippet.as_ref())
.map(|t| tokenize(t).len())
.collect();
if texts.is_empty() {
1.0
} else {
texts.iter().sum::<usize>() as f64 / texts.len() as f64
}
};
for cand in candidates.iter_mut() {
cand.modality_scores.clear();
if let (Some(qt), Some(snippet)) = (query_text, cand.text_snippet.as_deref()) {
let tf = self.compute_text_features(qt, snippet, avg_doc_len);
let text_score = (tf.bm25_score + tf.exact_match_bonus) * tf.length_penalty;
cand.modality_scores.push(ModalityScore::new(
"text",
text_score,
self.config.text_weight,
));
}
if let (Some(qe), Some(ce)) = (query_embedding, cand.embedding.as_deref()) {
let vf = Self::compute_vector_features(qe, ce)?;
cand.modality_scores.push(ModalityScore::new(
"vector",
vf.cosine_similarity,
self.config.vector_weight,
));
}
}
self.apply_fusion(&mut candidates)?;
candidates.sort_by(|a, b| {
b.final_score
.partial_cmp(&a.final_score)
.unwrap_or(std::cmp::Ordering::Equal)
});
if self.config.normalize_scores {
Self::normalize_scores(&mut candidates);
}
let original_ranks: Vec<(String, usize)> = candidates
.iter()
.enumerate()
.map(|(i, c)| (c.id.clone(), i + 1))
.collect();
candidates.retain(|c| c.final_score >= self.config.min_score_threshold);
if self.config.top_k > 0 && candidates.len() > self.config.top_k {
candidates.truncate(self.config.top_k);
}
for (i, c) in candidates.iter_mut().enumerate() {
c.rank = i + 1;
}
let total = candidates.len() as u64;
let displacement: f64 = candidates
.iter()
.map(|c| {
original_ranks
.iter()
.find(|(id, _)| id == &c.id)
.map(|(_, orig)| (c.rank as i64 - *orig as i64).unsigned_abs() as f64)
.unwrap_or(0.0)
})
.sum::<f64>()
/ total.max(1) as f64;
self.stats.candidates_reranked += total;
self.stats.fusion_calls += 1;
if self.stats.fusion_calls == 1 {
self.stats.avg_rank_displacement = displacement;
} else {
let n = self.stats.fusion_calls as f64;
self.stats.avg_rank_displacement =
(self.stats.avg_rank_displacement * (n - 1.0) + displacement) / n;
}
for c in &candidates {
for ms in &c.modality_scores {
if !self.stats.modalities_used.contains(&ms.modality) {
self.stats.modalities_used.push(ms.modality.clone());
}
}
}
Ok(candidates)
}
fn apply_fusion(&self, candidates: &mut [RerankerCandidate]) -> Result<(), RerankerError> {
match &self.config.fusion_strategy {
CmrFusionStrategy::LinearCombination(pairs) => {
let weight_map: HashMap<&str, f64> =
pairs.iter().map(|(k, v)| (k.as_str(), *v)).collect();
for cand in candidates.iter_mut() {
let score: f64 = cand
.modality_scores
.iter()
.map(|ms| {
let w = weight_map
.get(ms.modality.as_str())
.copied()
.unwrap_or(ms.weight);
w * ms.raw_score
})
.sum();
cand.final_score = score;
}
}
CmrFusionStrategy::ReciprocalRankFusion(k) => {
let k = *k;
let modality_names: Vec<String> = {
let mut names: Vec<String> = Vec::new();
for c in candidates.iter() {
for ms in &c.modality_scores {
if !names.contains(&ms.modality) {
names.push(ms.modality.clone());
}
}
}
names
};
let rank_lists: Vec<Vec<String>> = modality_names
.iter()
.map(|m| {
let mut scored: Vec<(String, f64)> = candidates
.iter()
.filter_map(|c| {
c.modality_scores
.iter()
.find(|ms| &ms.modality == m)
.map(|ms| (c.id.clone(), ms.raw_score))
})
.collect();
scored.sort_by(|a, b| {
b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
});
scored.into_iter().map(|(id, _)| id).collect()
})
.collect();
let rrf_scores = Self::reciprocal_rank_fusion(rank_lists, k);
let score_map: HashMap<&str, f64> =
rrf_scores.iter().map(|(id, s)| (id.as_str(), *s)).collect();
for cand in candidates.iter_mut() {
cand.final_score = score_map.get(cand.id.as_str()).copied().unwrap_or(0.0);
}
}
CmrFusionStrategy::Borda => {
let n = candidates.len();
let modality_names: Vec<String> = {
let mut names: Vec<String> = Vec::new();
for c in candidates.iter() {
for ms in &c.modality_scores {
if !names.contains(&ms.modality) {
names.push(ms.modality.clone());
}
}
}
names
};
let mut borda_totals: HashMap<String, f64> =
candidates.iter().map(|c| (c.id.clone(), 0.0)).collect();
for m in &modality_names {
let mut scored: Vec<(String, f64)> = candidates
.iter()
.filter_map(|c| {
c.modality_scores
.iter()
.find(|ms| &ms.modality == m)
.map(|ms| (c.id.clone(), ms.raw_score))
})
.collect();
scored
.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
for (rank_zero, (id, _)) in scored.iter().enumerate() {
let points = (n - rank_zero) as f64;
if let Some(total) = borda_totals.get_mut(id) {
*total += points;
}
}
}
for cand in candidates.iter_mut() {
cand.final_score = borda_totals.get(&cand.id).copied().unwrap_or(0.0);
}
}
CmrFusionStrategy::MaxScore => {
for cand in candidates.iter_mut() {
cand.final_score = cand
.modality_scores
.iter()
.map(|ms| ms.raw_score)
.fold(f64::NEG_INFINITY, f64::max);
if cand.final_score.is_infinite() {
cand.final_score = 0.0;
}
}
}
CmrFusionStrategy::LearnedWeights(weights) => {
for cand in candidates.iter_mut() {
let score: f64 = cand
.modality_scores
.iter()
.enumerate()
.map(|(i, ms)| {
let w = weights.get(i).copied().unwrap_or(1.0);
w * ms.raw_score
})
.sum();
cand.final_score = score;
}
}
}
Ok(())
}
fn normalize_scores(candidates: &mut [RerankerCandidate]) {
if candidates.is_empty() {
return;
}
let min_s = candidates
.iter()
.map(|c| c.final_score)
.fold(f64::INFINITY, f64::min);
let max_s = candidates
.iter()
.map(|c| c.final_score)
.fold(f64::NEG_INFINITY, f64::max);
let range = max_s - min_s;
if range < f64::EPSILON {
for c in candidates.iter_mut() {
c.final_score = 1.0;
}
return;
}
for c in candidates.iter_mut() {
c.final_score = (c.final_score - min_s) / range;
}
let modality_names: Vec<String> = {
let mut names: Vec<String> = Vec::new();
for c in candidates.iter() {
for ms in &c.modality_scores {
if !names.contains(&ms.modality) {
names.push(ms.modality.clone());
}
}
}
names
};
for m in &modality_names {
let min_r = candidates
.iter()
.flat_map(|c| c.modality_scores.iter())
.filter(|ms| &ms.modality == m)
.map(|ms| ms.raw_score)
.fold(f64::INFINITY, f64::min);
let max_r = candidates
.iter()
.flat_map(|c| c.modality_scores.iter())
.filter(|ms| &ms.modality == m)
.map(|ms| ms.raw_score)
.fold(f64::NEG_INFINITY, f64::max);
let r = max_r - min_r;
for c in candidates.iter_mut() {
for ms in c.modality_scores.iter_mut() {
if &ms.modality == m {
ms.normalized_score = if r < f64::EPSILON {
1.0
} else {
(ms.raw_score - min_r) / r
};
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_text_candidate(id: &str, text: &str) -> RerankerCandidate {
RerankerCandidate::new(id, Some(text), None)
}
fn make_vec_candidate(id: &str, embedding: Vec<f64>) -> RerankerCandidate {
RerankerCandidate::new(id, None, Some(embedding))
}
fn make_full_candidate(id: &str, text: &str, embedding: Vec<f64>) -> RerankerCandidate {
RerankerCandidate::new(id, Some(text), Some(embedding))
}
fn default_reranker() -> CrossModalReranker {
CrossModalReranker::new(RerankerConfig::default())
}
#[test]
fn test_tokenize_basic() {
let tokens = tokenize("Hello, World!");
assert_eq!(tokens, vec!["hello", "world"]);
}
#[test]
fn test_tokenize_empty() {
assert!(tokenize("").is_empty());
}
#[test]
fn test_tokenize_punctuation_stripped() {
let tokens = tokenize("rust, systems, programming.");
assert_eq!(tokens, vec!["rust", "systems", "programming"]);
}
#[test]
fn test_tokenize_lowercase() {
let tokens = tokenize("Rust PROGRAMMING");
assert!(tokens
.iter()
.all(|t| t.chars().all(|c| c.is_lowercase() || !c.is_alphabetic())));
}
#[test]
fn test_xorshift64_not_zero_after_seed() {
let mut state: u64 = 12345;
let v = xorshift64(&mut state);
assert_ne!(v, 0);
}
#[test]
fn test_xorshift_f64_in_range() {
let mut state: u64 = 99999;
for _ in 0..1000 {
let v = xorshift_f64(&mut state);
assert!((0.0..1.0).contains(&v), "value out of range: {v}");
}
}
#[test]
fn test_xorshift_f64_deterministic() {
let mut s1: u64 = 42;
let mut s2: u64 = 42;
assert_eq!(xorshift_f64(&mut s1), xorshift_f64(&mut s2));
}
#[test]
fn test_bm25_empty_query() {
let r = default_reranker();
let tf = r.compute_text_features("", "some text here", 4.0);
assert_eq!(tf.bm25_score, 0.0);
}
#[test]
fn test_bm25_empty_document() {
let r = default_reranker();
let tf = r.compute_text_features("rust", "", 4.0);
assert_eq!(tf.bm25_score, 0.0);
}
#[test]
fn test_bm25_term_present_vs_absent() {
let r = default_reranker();
let tf_present = r.compute_text_features("rust", "rust systems", 2.0);
let tf_absent = r.compute_text_features("rust", "python systems", 2.0);
assert!(tf_present.bm25_score > tf_absent.bm25_score);
}
#[test]
fn test_bm25_exact_match_bonus() {
let r = default_reranker();
let tf_exact =
r.compute_text_features("rust programming", "I love rust programming a lot", 5.0);
let tf_partial =
r.compute_text_features("rust programming", "I love rust and programming", 5.0);
assert!(
tf_exact.exact_match_bonus > tf_partial.exact_match_bonus,
"exact match should have bonus: exact={}, partial={}",
tf_exact.exact_match_bonus,
tf_partial.exact_match_bonus
);
}
#[test]
fn test_bm25_exact_match_bonus_value() {
let r = default_reranker();
let tf = r.compute_text_features("hello world", "hello world this is a test", 5.0);
assert!((tf.exact_match_bonus - 0.5).abs() < 1e-10);
}
#[test]
fn test_bm25_no_exact_match_bonus() {
let r = default_reranker();
let tf = r.compute_text_features("hello world", "goodbye everyone", 5.0);
assert_eq!(tf.exact_match_bonus, 0.0);
}
#[test]
fn test_bm25_length_penalty_short_doc() {
let r = default_reranker();
let tf = r.compute_text_features("a", "a", 100.0);
assert!((tf.length_penalty - 1.0).abs() < 1e-10);
}
#[test]
fn test_bm25_length_penalty_long_doc() {
let r = default_reranker();
let long_text = "word ".repeat(100);
let tf = r.compute_text_features("word", long_text.trim(), 10.0);
assert!(tf.length_penalty < 1.0);
}
#[test]
fn test_bm25_term_frequency_populated() {
let r = default_reranker();
let tf = r.compute_text_features("rust python", "rust is great", 3.0);
assert!(!tf.term_frequency.is_empty());
}
#[test]
fn test_bm25_zero_avg_doc_len_fallback() {
let r = default_reranker();
let tf = r.compute_text_features("hello", "hello world", 0.0);
assert!(tf.bm25_score.is_finite());
}
#[test]
fn test_bm25_custom_k1_b() {
let config = RerankerConfig {
bm25_k1: 2.0,
bm25_b: 0.5,
..Default::default()
};
let r = CrossModalReranker::new(config);
let tf = r.compute_text_features("rust", "rust systems rust", 3.0);
assert!(tf.bm25_score > 0.0);
}
#[test]
fn test_vector_features_identical() {
let v = vec![1.0, 0.0, 0.0];
let vf = CrossModalReranker::compute_vector_features(&v, &v)
.expect("test: identical vectors should compute without error");
assert!((vf.cosine_similarity - 1.0).abs() < 1e-10);
assert!(vf.l2_distance.abs() < 1e-10);
assert!((vf.euclidean_normalized - 1.0).abs() < 1e-10);
}
#[test]
fn test_vector_features_orthogonal() {
let q = vec![1.0, 0.0];
let c = vec![0.0, 1.0];
let vf = CrossModalReranker::compute_vector_features(&q, &c)
.expect("test: orthogonal vectors should compute without error");
assert!(vf.cosine_similarity.abs() < 1e-10);
}
#[test]
fn test_vector_features_opposite() {
let q = vec![1.0, 0.0];
let c = vec![-1.0, 0.0];
let vf = CrossModalReranker::compute_vector_features(&q, &c)
.expect("test: opposite vectors should compute without error");
assert!((vf.cosine_similarity + 1.0).abs() < 1e-10);
}
#[test]
fn test_vector_features_dimension_mismatch() {
let q = vec![1.0, 2.0, 3.0];
let c = vec![1.0, 2.0];
let err = CrossModalReranker::compute_vector_features(&q, &c)
.expect_err("test: dimension mismatch should return error");
assert_eq!(
err,
RerankerError::IncompatibleDimensions {
expected: 3,
got: 2
}
);
}
#[test]
fn test_vector_features_zero_vector() {
let q = vec![0.0, 0.0];
let c = vec![1.0, 0.0];
let vf = CrossModalReranker::compute_vector_features(&q, &c)
.expect("test: zero query vector should compute without error");
assert_eq!(vf.cosine_similarity, 0.0);
}
#[test]
fn test_vector_features_dot_product() {
let q = vec![1.0, 2.0, 3.0];
let c = vec![4.0, 5.0, 6.0];
let vf = CrossModalReranker::compute_vector_features(&q, &c)
.expect("test: dot product computation should succeed");
assert!((vf.dot_product - 32.0).abs() < 1e-10);
}
#[test]
fn test_vector_features_l2_distance() {
let q = vec![0.0, 0.0];
let c = vec![3.0, 4.0];
let vf = CrossModalReranker::compute_vector_features(&q, &c)
.expect("test: L2 distance computation should succeed");
assert!((vf.l2_distance - 5.0).abs() < 1e-10);
}
#[test]
fn test_vector_features_euclidean_normalized_bounded() {
let mut state: u64 = 777;
let q: Vec<f64> = (0..8).map(|_| xorshift_f64(&mut state)).collect();
let c: Vec<f64> = (0..8).map(|_| xorshift_f64(&mut state)).collect();
let vf = CrossModalReranker::compute_vector_features(&q, &c)
.expect("test: euclidean_normalized computation should succeed");
assert!(vf.euclidean_normalized > 0.0);
assert!(vf.euclidean_normalized <= 1.0);
}
#[test]
fn test_rrf_single_list() {
let lists = vec![vec!["a".to_string(), "b".to_string(), "c".to_string()]];
let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
let a_score = scores
.iter()
.find(|(id, _)| id == "a")
.expect("test: 'a' must be in RRF scores")
.1;
let b_score = scores
.iter()
.find(|(id, _)| id == "b")
.expect("test: 'b' must be in RRF scores")
.1;
assert!(a_score > b_score);
}
#[test]
fn test_rrf_two_lists_consensus() {
let lists = vec![
vec!["a".to_string(), "b".to_string()],
vec!["a".to_string(), "b".to_string()],
];
let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
let a = scores
.iter()
.find(|(id, _)| id == "a")
.expect("test: 'a' must be in RRF scores")
.1;
let b = scores
.iter()
.find(|(id, _)| id == "b")
.expect("test: 'b' must be in RRF scores")
.1;
assert!(a > b);
}
#[test]
fn test_rrf_rank_disagreement() {
let lists = vec![
vec!["a".to_string(), "b".to_string()],
vec!["b".to_string(), "a".to_string()],
];
let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
let a = scores
.iter()
.find(|(id, _)| id == "a")
.expect("test: 'a' must be in RRF scores")
.1;
let b = scores
.iter()
.find(|(id, _)| id == "b")
.expect("test: 'b' must be in RRF scores")
.1;
assert!((a - b).abs() < 1e-10, "a={a}, b={b}");
}
#[test]
fn test_rrf_custom_k() {
let lists = vec![vec!["x".to_string()]];
let s1 = CrossModalReranker::reciprocal_rank_fusion(lists.clone(), 10.0);
let s2 = CrossModalReranker::reciprocal_rank_fusion(lists, 100.0);
let v1 = s1[0].1;
let v2 = s2[0].1;
assert!(v1 > v2);
}
#[test]
fn test_rrf_empty_lists() {
let scores = CrossModalReranker::reciprocal_rank_fusion(vec![], 60.0);
assert!(scores.is_empty());
}
#[test]
fn test_rrf_sorted_descending() {
let lists = vec![
vec!["c".to_string(), "b".to_string(), "a".to_string()],
vec!["a".to_string(), "b".to_string(), "c".to_string()],
];
let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
for w in scores.windows(2) {
assert!(w[0].1 >= w[1].1);
}
}
#[test]
fn test_text_only_rerank_ordering() {
let mut r = default_reranker();
let candidates = vec![
make_text_candidate("doc1", "python machine learning"),
make_text_candidate("doc2", "rust systems programming rust"),
];
let results = r
.rerank(candidates, Some("rust"), None)
.expect("test: rerank should succeed");
assert_eq!(results[0].id, "doc2");
}
#[test]
fn test_text_only_rerank_ranks_assigned() {
let mut r = default_reranker();
let candidates = vec![
make_text_candidate("a", "foo"),
make_text_candidate("b", "foo bar"),
make_text_candidate("c", "foo bar baz"),
];
let results = r
.rerank(candidates, Some("foo"), None)
.expect("test: rerank should succeed");
for (i, res) in results.iter().enumerate() {
assert_eq!(res.rank, i + 1);
}
}
#[test]
fn test_text_only_empty_query_still_returns() {
let mut r = default_reranker();
let candidates = vec![make_text_candidate("a", "hello world")];
let results = r
.rerank(candidates, Some(""), None)
.expect("test: rerank should succeed");
assert_eq!(results.len(), 1);
}
#[test]
fn test_vector_only_rerank_ordering() {
let mut r = default_reranker();
let query = vec![1.0_f64, 0.0];
let close = make_vec_candidate("close", vec![0.99, 0.14]);
let far = make_vec_candidate("far", vec![0.0, 1.0]);
let results = r
.rerank(vec![far, close], None, Some(&query))
.expect("test: rerank should succeed");
assert_eq!(results[0].id, "close");
}
#[test]
fn test_vector_only_rerank_scores_finite() {
let mut r = default_reranker();
let mut state: u64 = 1234;
let q: Vec<f64> = (0..16).map(|_| xorshift_f64(&mut state)).collect();
let candidates: Vec<RerankerCandidate> = (0..5)
.map(|i| {
let emb: Vec<f64> = (0..16).map(|_| xorshift_f64(&mut state)).collect();
make_vec_candidate(&format!("doc{i}"), emb)
})
.collect();
let results = r
.rerank(candidates, None, Some(&q))
.expect("test: rerank should succeed");
for res in &results {
assert!(res.final_score.is_finite());
}
}
#[test]
fn test_vector_only_dimension_mismatch_error() {
let mut r = default_reranker();
let q = vec![1.0, 2.0, 3.0];
let cand = make_vec_candidate("bad", vec![1.0, 2.0]);
let err = r
.rerank(vec![cand], None, Some(&q))
.expect_err("test: rerank should return error for dimension mismatch");
assert!(matches!(err, RerankerError::IncompatibleDimensions { .. }));
}
#[test]
fn test_cross_modal_fusion_both_modalities_present() {
let mut r = default_reranker();
let q_text = "rust programming";
let q_emb = vec![1.0_f64, 0.0];
let candidates = vec![
make_full_candidate("doc1", "rust programming language", vec![0.99, 0.14]),
make_full_candidate("doc2", "python data science", vec![0.0, 1.0]),
];
let results = r
.rerank(candidates, Some(q_text), Some(&q_emb))
.expect("test: rerank should succeed");
assert_eq!(results[0].id, "doc1");
}
#[test]
fn test_cross_modal_modality_scores_populated() {
let mut r = default_reranker();
let candidates = vec![make_full_candidate("doc1", "hello world", vec![1.0, 0.0])];
let q_emb = vec![1.0, 0.0];
let results = r
.rerank(candidates, Some("hello"), Some(&q_emb))
.expect("test: rerank should succeed");
assert!(!results[0].modality_scores.is_empty());
let has_text = results[0]
.modality_scores
.iter()
.any(|ms| ms.modality == "text");
let has_vec = results[0]
.modality_scores
.iter()
.any(|ms| ms.modality == "vector");
assert!(has_text);
assert!(has_vec);
}
#[test]
fn test_linear_combination_weights() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::LinearCombination(vec![
("text".to_string(), 0.9),
("vector".to_string(), 0.1),
]),
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_full_candidate("doc1", "rust is great", vec![0.0, 1.0]),
make_full_candidate("doc2", "python is fine", vec![1.0, 0.0]),
];
let q_emb = vec![1.0_f64, 0.0];
let results = r
.rerank(candidates, Some("rust"), Some(&q_emb))
.expect("test: rerank should succeed");
assert_eq!(results[0].id, "doc1");
}
#[test]
fn test_rrf_fusion_strategy() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::ReciprocalRankFusion(60.0),
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_full_candidate("doc1", "rust rust rust", vec![0.9, 0.0]),
make_full_candidate("doc2", "python", vec![0.1, 0.0]),
];
let q_emb = vec![1.0_f64, 0.0];
let results = r
.rerank(candidates, Some("rust"), Some(&q_emb))
.expect("test: rerank should succeed");
assert!(!results.is_empty());
}
#[test]
fn test_borda_fusion_strategy() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::Borda,
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_full_candidate("doc1", "rust rust", vec![0.8, 0.2]),
make_full_candidate("doc2", "python", vec![0.2, 0.8]),
];
let q_emb = vec![1.0_f64, 0.0];
let results = r
.rerank(candidates, Some("rust"), Some(&q_emb))
.expect("test: rerank should succeed");
for res in &results {
assert!(res.final_score >= 0.0);
}
}
#[test]
fn test_borda_scores_non_negative() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::Borda,
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = (0..5)
.map(|i| make_text_candidate(&format!("d{i}"), &"word ".repeat(i + 1)))
.collect();
let results = r
.rerank(candidates, Some("word"), None)
.expect("test: rerank should succeed");
for res in &results {
assert!(res.final_score >= 0.0);
}
}
#[test]
fn test_max_score_fusion() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::MaxScore,
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_full_candidate("doc1", "rust rust rust rust", vec![0.2, 0.0]),
make_full_candidate("doc2", "python", vec![0.99, 0.0]),
];
let q_emb = vec![1.0_f64, 0.0];
let results = r
.rerank(candidates, Some("rust"), Some(&q_emb))
.expect("test: rerank should succeed");
assert!(!results.is_empty());
for res in &results {
assert!(res.final_score.is_finite());
}
}
#[test]
fn test_learned_weights_fusion() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::LearnedWeights(vec![2.0, 1.0]),
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_full_candidate("doc1", "rust programming", vec![0.5, 0.0]),
make_full_candidate("doc2", "java programming", vec![0.8, 0.0]),
];
let q_emb = vec![1.0_f64, 0.0];
let results = r
.rerank(candidates, Some("rust"), Some(&q_emb))
.expect("test: rerank should succeed");
assert!(!results.is_empty());
}
#[test]
fn test_normalize_scores_in_range() {
let config = RerankerConfig {
normalize_scores: true,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let mut state: u64 = 54321;
let q: Vec<f64> = (0..4).map(|_| xorshift_f64(&mut state)).collect();
let candidates: Vec<RerankerCandidate> = (0..8)
.map(|i| {
let emb: Vec<f64> = (0..4).map(|_| xorshift_f64(&mut state)).collect();
make_full_candidate(&format!("d{i}"), "some text here", emb)
})
.collect();
let results = r
.rerank(candidates, Some("text"), Some(&q))
.expect("test: rerank should succeed");
for res in &results {
assert!(
(0.0..=1.0).contains(&res.final_score),
"score={}",
res.final_score
);
}
}
#[test]
fn test_normalize_scores_disabled() {
let config = RerankerConfig {
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_text_candidate("d1", "hello world hello"),
make_text_candidate("d2", "foo bar"),
];
let results = r
.rerank(candidates, Some("hello"), None)
.expect("test: rerank should succeed");
for res in &results {
assert!(res.final_score.is_finite());
}
}
#[test]
fn test_normalize_all_equal_scores() {
let config = RerankerConfig {
normalize_scores: true,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_vec_candidate("d1", vec![1.0, 0.0]),
make_vec_candidate("d2", vec![1.0, 0.0]),
];
let q = vec![1.0, 0.0];
let results = r
.rerank(candidates, None, Some(&q))
.expect("test: rerank should succeed");
for res in &results {
assert!((0.0..=1.0).contains(&res.final_score));
}
}
#[test]
fn test_top_k_limits_results() {
let config = RerankerConfig {
top_k: 3,
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = (0..10)
.map(|i| make_text_candidate(&format!("d{i}"), &format!("word {i}")))
.collect();
let results = r
.rerank(candidates, Some("word"), None)
.expect("test: rerank should succeed");
assert!(results.len() <= 3);
}
#[test]
fn test_top_k_zero_returns_all() {
let config = RerankerConfig {
top_k: 0,
normalize_scores: false,
min_score_threshold: 0.0,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = (0..5)
.map(|i| make_text_candidate(&format!("d{i}"), &format!("word {i}")))
.collect();
let results = r
.rerank(candidates, Some("word"), None)
.expect("test: rerank should succeed");
assert_eq!(results.len(), 5);
}
#[test]
fn test_min_score_threshold_filters_low() {
let config = RerankerConfig {
normalize_scores: true,
min_score_threshold: 0.5,
top_k: 0,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_text_candidate("high", "rust rust rust rust"),
make_text_candidate("low", "java"),
];
let results = r
.rerank(candidates, Some("rust"), None)
.expect("test: rerank should succeed");
for res in &results {
assert!(
res.final_score >= 0.5,
"score below threshold: {}",
res.final_score
);
}
}
#[test]
fn test_min_score_threshold_zero_keeps_all() {
let config = RerankerConfig {
normalize_scores: false,
min_score_threshold: 0.0,
top_k: 0,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = (0..4)
.map(|i| make_text_candidate(&format!("d{i}"), &format!("text {i}")))
.collect();
let results = r
.rerank(candidates, Some("text"), None)
.expect("test: rerank should succeed");
assert_eq!(results.len(), 4);
}
#[test]
fn test_error_no_candidates() {
let mut r = default_reranker();
let err = r
.rerank(vec![], Some("query"), None)
.expect_err("test: rerank should return error for empty candidates");
assert_eq!(err, RerankerError::NoCandidates);
}
#[test]
fn test_error_incompatible_dimensions() {
let mut r = default_reranker();
let cand = make_vec_candidate("d1", vec![1.0, 2.0]);
let q = vec![1.0, 2.0, 3.0];
let err = r
.rerank(vec![cand], None, Some(&q))
.expect_err("test: rerank should return error for incompatible dimensions");
assert!(matches!(err, RerankerError::IncompatibleDimensions { .. }));
}
#[test]
fn test_error_invalid_weight_negative() {
let config = RerankerConfig {
text_weight: -1.0,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let cand = make_text_candidate("d1", "hello");
let err = r
.rerank(vec![cand], Some("hello"), None)
.expect_err("test: rerank should return error for negative text_weight");
assert!(matches!(err, RerankerError::InvalidWeight(_)));
}
#[test]
fn test_error_invalid_linear_weight() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::LinearCombination(vec![("text".to_string(), -0.1)]),
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let cand = make_text_candidate("d1", "hello");
let err = r
.rerank(vec![cand], Some("hello"), None)
.expect_err("test: rerank should return error for negative linear combination weight");
assert!(matches!(err, RerankerError::InvalidWeight(_)));
}
#[test]
fn test_error_invalid_rrf_k() {
let config = RerankerConfig {
fusion_strategy: CmrFusionStrategy::ReciprocalRankFusion(-1.0),
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let cand = make_text_candidate("d1", "hello");
let err = r
.rerank(vec![cand], Some("hello"), None)
.expect_err("test: rerank should return error for invalid RRF k value");
assert!(matches!(err, RerankerError::ConfigurationError(_)));
}
#[test]
fn test_error_display() {
let e = RerankerError::NoCandidates;
assert!(!format!("{e}").is_empty());
let e2 = RerankerError::IncompatibleDimensions {
expected: 4,
got: 3,
};
assert!(format!("{e2}").contains("4"));
}
#[test]
fn test_stats_fusion_calls_incremented() {
let mut r = default_reranker();
assert_eq!(r.stats().fusion_calls, 0);
let _ = r.rerank(vec![make_text_candidate("a", "hi")], Some("hi"), None);
assert_eq!(r.stats().fusion_calls, 1);
let _ = r.rerank(vec![make_text_candidate("b", "bye")], Some("bye"), None);
assert_eq!(r.stats().fusion_calls, 2);
}
#[test]
fn test_stats_candidates_reranked_accumulates() {
let mut r = default_reranker();
let c1 = vec![make_text_candidate("a", "a"), make_text_candidate("b", "b")];
let c2 = vec![make_text_candidate("c", "c")];
let _ = r.rerank(c1, Some("q"), None);
let _ = r.rerank(c2, Some("q"), None);
assert_eq!(r.stats().candidates_reranked, 3);
}
#[test]
fn test_stats_modalities_tracked() {
let mut r = default_reranker();
let cands = vec![make_full_candidate("d1", "hello", vec![1.0, 0.0])];
let q_emb = vec![1.0, 0.0];
let _ = r.rerank(cands, Some("hello"), Some(&q_emb));
let stats = r.stats();
assert!(stats.modalities_used.contains(&"text".to_string()));
assert!(stats.modalities_used.contains(&"vector".to_string()));
}
#[test]
fn test_update_config() {
let mut r = default_reranker();
let new_cfg = RerankerConfig {
top_k: 5,
..Default::default()
};
r.update_config(new_cfg.clone());
assert_eq!(r.config.top_k, 5);
}
#[test]
fn test_modality_score_new() {
let ms = ModalityScore::new("text", 0.8, 0.4);
assert_eq!(ms.modality, "text");
assert!((ms.raw_score - 0.8).abs() < 1e-10);
assert!((ms.normalized_score - 0.8).abs() < 1e-10);
assert!((ms.weight - 0.4).abs() < 1e-10);
}
#[test]
fn test_single_candidate_gets_rank_one() {
let mut r = default_reranker();
let cand = make_text_candidate("solo", "only document");
let results = r
.rerank(vec![cand], Some("document"), None)
.expect("test: rerank should succeed");
assert_eq!(results[0].rank, 1);
}
#[test]
fn test_single_candidate_score_normalised_to_one() {
let config = RerankerConfig {
normalize_scores: true,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let cand = make_text_candidate("solo", "only document");
let results = r
.rerank(vec![cand], Some("document"), None)
.expect("test: rerank should succeed");
assert!((results[0].final_score - 1.0).abs() < 1e-10);
}
#[test]
fn test_candidates_with_no_matching_modality_get_zero_score() {
let config = RerankerConfig {
normalize_scores: false,
..Default::default()
};
let mut r = CrossModalReranker::new(config);
let candidates = vec![
make_text_candidate("d1", "hello"),
make_text_candidate("d2", "world"),
];
let results = r
.rerank(candidates, None, None)
.expect("test: rerank should succeed");
for res in &results {
assert_eq!(res.final_score, 0.0);
}
}
#[test]
fn test_reranker_candidate_new() {
let c = RerankerCandidate::new("id", Some("text"), Some(vec![1.0]));
assert_eq!(c.id, "id");
assert_eq!(c.text_snippet.as_deref(), Some("text"));
assert_eq!(c.embedding, Some(vec![1.0]));
assert!(c.modality_scores.is_empty());
assert_eq!(c.rank, 0);
}
#[test]
fn test_large_candidate_set_no_panic() {
let mut state: u64 = 314159;
let mut r = default_reranker();
let q: Vec<f64> = (0..32).map(|_| xorshift_f64(&mut state)).collect();
let candidates: Vec<RerankerCandidate> = (0..200)
.map(|i| {
let emb: Vec<f64> = (0..32).map(|_| xorshift_f64(&mut state)).collect();
make_full_candidate(&format!("d{i}"), "some query terms here", emb)
})
.collect();
let results = r
.rerank(candidates, Some("query terms"), Some(&q))
.expect("test: rerank should succeed");
assert!(results.len() <= 100); }
}