hermes_core/query/vector/

combiner.rs

//! Multi-value score combination strategies for vector search

/// Strategy for combining scores when a document has multiple values for the same field
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MultiValueCombiner {
    /// Sum all scores (accumulates dot product contributions)
    Sum,
    /// Take the maximum score
    Max,
    /// Take the average score
    Avg,
    /// Softmax-weighted smooth maximum (default)
    /// `score = Σ softmax(t * sᵢ) * sᵢ`
    /// Higher temperature → closer to max; lower → closer to mean.
    ///
    /// Deliberately NOT the raw `(1/t)·log(Σ exp(t·sᵢ))`: that form adds
    /// `ln(n)/t` per document, so chunk *count* outranked chunk quality —
    /// a 300-chunk compendium of mediocre matches beat every focused paper
    /// (score 0.55 + ln(300)/1.5 ≈ 4.3 vs 0.72 + ln(3)/1.5 ≈ 1.4). The
    /// softmax weighting is count-invariant (n identical scores combine to
    /// that score), bounded by the max, and still tracks a dominant score
    /// at any scale.
    LogSumExp {
        /// Temperature parameter (default: 1.5)
        temperature: f32,
    },
    /// Weighted Top-K: weight top scores with exponential decay
    /// `score = Σ wᵢ * sorted_scores[i]` where `wᵢ = decay^i`
    WeightedTopK {
        /// Number of top scores to consider (default: 5)
        k: usize,
        /// Decay factor per rank (default: 0.7)
        decay: f32,
    },
}

impl Default for MultiValueCombiner {
    fn default() -> Self {
        // LogSumExp with temperature 1.5 provides good balance between
        // max (best relevance) and sum (saturation from multiple matches)
        MultiValueCombiner::LogSumExp { temperature: 1.5 }
    }
}

impl MultiValueCombiner {
    pub(crate) fn validate(self) -> Result<(), String> {
        match self {
            Self::LogSumExp { temperature } if !temperature.is_finite() || temperature <= 0.0 => {
                Err(format!(
                    "LogSumExp temperature must be finite and greater than zero, got {temperature}"
                ))
            }
            Self::WeightedTopK { k: 0, .. } => {
                Err("WeightedTopK k must be greater than zero".to_string())
            }
            Self::WeightedTopK { decay, .. }
                if !decay.is_finite() || !(0.0..=1.0).contains(&decay) =>
            {
                Err(format!(
                    "WeightedTopK decay must be finite and in [0, 1], got {decay}"
                ))
            }
            _ => Ok(()),
        }
    }

    /// Create LogSumExp combiner with default temperature (1.5)
    pub fn log_sum_exp() -> Self {
        Self::LogSumExp { temperature: 1.5 }
    }

    /// Create LogSumExp combiner with custom temperature
    pub fn log_sum_exp_with_temperature(temperature: f32) -> Self {
        Self::LogSumExp { temperature }
    }

    /// Create WeightedTopK combiner with defaults (k=5, decay=0.7)
    pub fn weighted_top_k() -> Self {
        Self::WeightedTopK { k: 5, decay: 0.7 }
    }

    /// Create WeightedTopK combiner with custom parameters
    pub fn weighted_top_k_with_params(k: usize, decay: f32) -> Self {
        Self::WeightedTopK { k, decay }
    }

    /// Combine multiple scores into a single score
    pub fn combine(&self, scores: &[(u32, f32)]) -> f32 {
        if scores.is_empty() {
            return 0.0;
        }

        match self {
            MultiValueCombiner::Sum => scores.iter().map(|(_, s)| s).sum(),
            MultiValueCombiner::Max => scores
                .iter()
                .map(|(_, s)| *s)
                .max_by(|a, b| a.total_cmp(b))
                .unwrap_or(0.0),
            MultiValueCombiner::Avg => {
                let sum: f32 = scores.iter().map(|(_, s)| s).sum();
                sum / scores.len() as f32
            }
            MultiValueCombiner::LogSumExp { temperature } => {
                // Softmax-weighted average, numerically stabilized by
                // subtracting the max before exponentiation. The max's own
                // weight is exp(0) = 1, so the denominator is never zero.
                let t = *temperature;
                let max_score = scores
                    .iter()
                    .map(|(_, s)| *s)
                    .max_by(|a, b| a.total_cmp(b))
                    .unwrap_or(0.0);

                let mut weight_sum = 0.0f32;
                let mut weighted = 0.0f32;
                for (_, s) in scores {
                    let weight = (t * (s - max_score)).exp();
                    weight_sum += weight;
                    weighted += weight * s;
                }
                weighted / weight_sum
            }
            MultiValueCombiner::WeightedTopK { k, decay } => {
                // Sort scores descending and take top k
                let mut sorted: Vec<f32> = scores.iter().map(|(_, s)| *s).collect();
                sorted.sort_unstable_by(|a, b| b.total_cmp(a));
                sorted.truncate(*k);

                // Apply exponential decay weights
                let mut weight = 1.0f32;
                let mut weighted_sum = 0.0f32;
                let mut weight_total = 0.0f32;

                for score in sorted {
                    weighted_sum += weight * score;
                    weight_total += weight;
                    weight *= decay;
                }

                if weight_total > 0.0 {
                    weighted_sum / weight_total
                } else {
                    0.0
                }
            }
        }
    }
}

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

    #[test]
    fn test_combiner_sum() {
        let scores = vec![(0, 1.0), (1, 2.0), (2, 3.0)];
        let combiner = MultiValueCombiner::Sum;
        assert!((combiner.combine(&scores) - 6.0).abs() < 1e-6);
    }

    #[test]
    fn test_combiner_max() {
        let scores = vec![(0, 1.0), (1, 3.0), (2, 2.0)];
        let combiner = MultiValueCombiner::Max;
        assert!((combiner.combine(&scores) - 3.0).abs() < 1e-6);
    }

    #[test]
    fn test_combiner_avg() {
        let scores = vec![(0, 1.0), (1, 2.0), (2, 3.0)];
        let combiner = MultiValueCombiner::Avg;
        assert!((combiner.combine(&scores) - 2.0).abs() < 1e-6);
    }

    #[test]
    fn test_combiner_log_sum_exp() {
        let scores = vec![(0, 1.0), (1, 2.0), (2, 3.0)];
        let combiner = MultiValueCombiner::log_sum_exp();
        let result = combiner.combine(&scores);
        // A smooth maximum lives between the mean and the max, weighted
        // toward the max.
        assert!(result > 2.0, "must exceed the mean, got {result}");
        assert!(result <= 3.0, "must never exceed the max, got {result}");
    }

    /// The production incident this pins: a 300-chunk compendium whose chunks
    /// all score ~0.5 must not outrank a 3-chunk paper whose chunks score
    /// ~0.7. The additive `ln(n)/t` count term of the raw log-sum-exp did
    /// exactly that (0.55 + ln(300)/1.5 ≈ 4.3 vs 0.72 + ln(3)/1.5 ≈ 1.4),
    /// which buried every relevant result for "off-label aripiprazole usage"
    /// under generic long documents.
    #[test]
    fn log_sum_exp_is_count_invariant_and_bounded_by_max() {
        let combiner = MultiValueCombiner::log_sum_exp();

        // n identical scores combine to that score, regardless of n.
        let identical: Vec<(u32, f32)> = (0..300).map(|i| (i, 0.7)).collect();
        let combined = combiner.combine(&identical);
        assert!(
            (combined - 0.7).abs() < 1e-3,
            "300 identical 0.7 chunks must combine to 0.7, got {combined}"
        );

        // Many mediocre chunks never outrank a few strong ones.
        let mut compendium: Vec<(u32, f32)> = (0..300).map(|i| (i, 0.5)).collect();
        compendium.push((300, 0.55));
        let paper = vec![(0, 0.72), (1, 0.70), (2, 0.65)];
        let compendium_score = combiner.combine(&compendium);
        let paper_score = combiner.combine(&paper);
        assert!(
            paper_score > compendium_score,
            "3 strong chunks ({paper_score}) must beat 301 mediocre ones ({compendium_score})"
        );
    }

    /// The reason the fix is a softmax-weighted average rather than
    /// `LSE - ln(n)/t`: subtracting the count term turns the combiner into a
    /// near-average in the peaked regime, collapsing a document whose single
    /// chunk scores 15 among 299 zeros to ~6.4. The softmax weighting keeps
    /// it at the dominant score.
    #[test]
    fn log_sum_exp_tracks_a_dominant_score_at_sparse_scale() {
        let combiner = MultiValueCombiner::log_sum_exp_with_temperature(0.7);
        let mut scores: Vec<(u32, f32)> = (0..299).map(|i| (i, 0.0)).collect();
        scores.push((299, 15.0));
        let combined = combiner.combine(&scores);
        assert!(
            (combined - 15.0).abs() < 0.3,
            "one dominant sparse chunk must keep its score, got {combined}"
        );
    }

    #[test]
    fn test_combiner_log_sum_exp_approaches_max_with_high_temp() {
        let scores = vec![(0, 1.0), (1, 5.0), (2, 2.0)];
        // High temperature should approach max
        let combiner = MultiValueCombiner::log_sum_exp_with_temperature(10.0);
        let result = combiner.combine(&scores);
        // Should be very close to max (5.0)
        assert!((result - 5.0).abs() < 0.5);
    }

    #[test]
    fn test_combiner_weighted_top_k() {
        let scores = vec![(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)];
        let combiner = MultiValueCombiner::weighted_top_k_with_params(3, 0.5);
        let result = combiner.combine(&scores);
        // Top 3: 5.0, 3.0, 1.0 with weights 1.0, 0.5, 0.25
        // weighted_sum = 5*1 + 3*0.5 + 1*0.25 = 6.75
        // weight_total = 1.75
        // result = 6.75 / 1.75 ≈ 3.857
        assert!((result - 3.857).abs() < 0.01);
    }

    #[test]
    fn test_combiner_weighted_top_k_less_than_k() {
        let scores = vec![(0, 2.0), (1, 1.0)];
        let combiner = MultiValueCombiner::weighted_top_k_with_params(5, 0.7);
        let result = combiner.combine(&scores);
        // Only 2 scores, weights 1.0 and 0.7
        // weighted_sum = 2*1 + 1*0.7 = 2.7
        // weight_total = 1.7
        // result = 2.7 / 1.7 ≈ 1.588
        assert!((result - 1.588).abs() < 0.01);
    }

    #[test]
    fn test_combiner_empty_scores() {
        let scores: Vec<(u32, f32)> = vec![];
        assert_eq!(MultiValueCombiner::Sum.combine(&scores), 0.0);
        assert_eq!(MultiValueCombiner::Max.combine(&scores), 0.0);
        assert_eq!(MultiValueCombiner::Avg.combine(&scores), 0.0);
        assert_eq!(MultiValueCombiner::log_sum_exp().combine(&scores), 0.0);
        assert_eq!(MultiValueCombiner::weighted_top_k().combine(&scores), 0.0);
    }

    #[test]
    fn test_combiner_single_score() {
        let scores = vec![(0, 5.0)];
        // All combiners should return 5.0 for a single score
        assert!((MultiValueCombiner::Sum.combine(&scores) - 5.0).abs() < 1e-6);
        assert!((MultiValueCombiner::Max.combine(&scores) - 5.0).abs() < 1e-6);
        assert!((MultiValueCombiner::Avg.combine(&scores) - 5.0).abs() < 1e-6);
        assert!((MultiValueCombiner::log_sum_exp().combine(&scores) - 5.0).abs() < 1e-6);
        assert!((MultiValueCombiner::weighted_top_k().combine(&scores) - 5.0).abs() < 1e-6);
    }

    #[test]
    fn test_default_combiner_is_log_sum_exp() {
        let combiner = MultiValueCombiner::default();
        match combiner {
            MultiValueCombiner::LogSumExp { temperature } => {
                assert!((temperature - 1.5).abs() < 1e-6);
            }
            _ => panic!("Default combiner should be LogSumExp"),
        }
    }
}