klieo-eval 1.0.0

//! Recall-quality metrics for klieo memory pipelines.
//!
//! Pure scoring — no I/O, no async, no dependency on the memory
//! traits. Callers run the recall against whatever pipeline they
//! want (Qdrant, sqlite-vec, hybrid GraphRAG, …), collect the
//! returned ids into a [`RecallSample`], and hand the slice to
//! [`score_recall`].
//!
//! # Metrics
//!
//! - **hit_rate** — fraction of queries with at least one expected
//!   id in the top-K result list.
//! - **mean reciprocal rank (MRR)** — mean of `1 / rank-of-first-hit`
//!   across queries; non-hit contributes 0.
//! - **mean precision@K** — `|expected ∩ top-K| / K`, averaged
//!   across queries.
//! - **mean recall@K** — `|expected ∩ top-K| / |expected|`,
//!   averaged across queries with at least one expected id.
//!   Queries with empty `expected` do not contribute to the mean.
//! - **mean NDCG@K** — normalised discounted cumulative gain over
//!   the top-K results, binary relevance (`1` if expected, `0`
//!   otherwise), averaged across queries with at least one
//!   expected id.
//!
//! # Example
//!
//! ```
//! use klieo_eval::{score_recall, RecallSample};
//!
//! let samples = vec![RecallSample::new(
//!     "ICT risk management",
//!     vec!["dora-art-5".into(), "dora-art-6".into()],
//!     vec![
//!         "dora-art-5".into(),
//!         "ai-act-art-6".into(),
//!         "dora-art-6".into(),
//!     ],
//! )];
//! let report = score_recall(&samples, 5);
//! assert!((report.hit_rate() - 1.0).abs() < 1e-9);
//! assert!(report.mean_reciprocal_rank() > 0.99);
//! ```

use serde::Serialize;

/// One scored query — the question, the ground-truth ids that
/// should land in the top-K, and the actual ranked id list
/// returned by the recall pipeline (top result first).
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct RecallSample {
    pub query: String,
    pub expected: Vec<String>,
    pub returned_ids: Vec<String>,
}

impl RecallSample {
    pub fn new(query: impl Into<String>, expected: Vec<String>, returned_ids: Vec<String>) -> Self {
        Self {
            query: query.into(),
            expected,
            returned_ids,
        }
    }
}

/// Per-query trace plus aggregate metrics computed across the
/// scored sample slice.
#[derive(Debug, Clone, Serialize)]
#[non_exhaustive]
pub struct RecallReport {
    pub k: usize,
    pub per_query: Vec<PerQueryScore>,
}

/// One row of the report — every metric for one query, plus the
/// rank of the first hit (1-based) or `None` when none of the
/// expected ids appeared in the top-K window. `expected_count`
/// is preserved so recall- and NDCG-shaped means can correctly
/// skip queries with empty `expected`.
#[derive(Debug, Clone, Serialize)]
#[non_exhaustive]
pub struct PerQueryScore {
    pub query: String,
    pub expected_count: usize,
    pub first_hit_rank: Option<usize>,
    pub precision_at_k: f64,
    pub recall_at_k: f64,
    pub ndcg_at_k: f64,
}

impl RecallReport {
    /// Fraction of queries where any expected id appeared in the
    /// top-K. `0.0` when there are no queries.
    pub fn hit_rate(&self) -> f64 {
        if self.per_query.is_empty() {
            return 0.0;
        }
        let hits = self
            .per_query
            .iter()
            .filter(|q| q.first_hit_rank.is_some())
            .count();
        hits as f64 / self.per_query.len() as f64
    }

    /// Mean of `1 / rank` of the first expected hit across all
    /// queries. Non-hits contribute `0`. `0.0` when there are no
    /// queries.
    pub fn mean_reciprocal_rank(&self) -> f64 {
        if self.per_query.is_empty() {
            return 0.0;
        }
        let sum: f64 = self
            .per_query
            .iter()
            .map(|q| {
                q.first_hit_rank
                    .map(|rank| 1.0 / rank as f64)
                    .unwrap_or(0.0)
            })
            .sum();
        sum / self.per_query.len() as f64
    }

    /// Mean precision@K across all queries. `0.0` when there are
    /// no queries.
    pub fn mean_precision_at_k(&self) -> f64 {
        mean(self.per_query.iter().map(|q| q.precision_at_k))
    }

    /// Mean recall@K across queries with at least one expected
    /// id. Queries with empty `expected` are excluded from the
    /// denominator — recall is undefined when there is nothing
    /// to retrieve. `0.0` when no query qualifies.
    pub fn mean_recall_at_k(&self) -> f64 {
        mean(
            self.per_query
                .iter()
                .filter(|q| q.expected_count > 0)
                .map(|q| q.recall_at_k),
        )
    }

    /// Mean NDCG@K across queries with at least one expected id.
    /// Same skip rule as [`Self::mean_recall_at_k`] — NDCG is
    /// undefined when ideal DCG is zero. `0.0` when no query
    /// qualifies.
    pub fn mean_ndcg_at_k(&self) -> f64 {
        mean(
            self.per_query
                .iter()
                .filter(|q| q.expected_count > 0)
                .map(|q| q.ndcg_at_k),
        )
    }
}

/// Score a batch of samples against a top-K window.
///
/// `k` is the window size used for every per-query metric — the
/// caller must truncate or pad `returned_ids` to honour their own
/// retrieval-limit semantics before scoring. The harness will
/// only ever inspect the first `k` entries of each
/// `returned_ids` list.
pub fn score_recall(samples: &[RecallSample], k: usize) -> RecallReport {
    let per_query = samples
        .iter()
        .map(|sample| score_sample(sample, k))
        .collect();
    RecallReport { k, per_query }
}

fn score_sample(sample: &RecallSample, k: usize) -> PerQueryScore {
    let window = sample.returned_ids.iter().take(k).collect::<Vec<_>>();
    let first_hit_rank = first_hit_rank(&window, &sample.expected);
    let precision_at_k = precision_at_k(&window, &sample.expected, k);
    let recall_at_k = recall_at_k(&window, &sample.expected);
    let ndcg_at_k = ndcg_at_k(&window, &sample.expected, k);
    PerQueryScore {
        query: sample.query.clone(),
        expected_count: sample.expected.len(),
        first_hit_rank,
        precision_at_k,
        recall_at_k,
        ndcg_at_k,
    }
}

fn first_hit_rank(window: &[&String], expected: &[String]) -> Option<usize> {
    window
        .iter()
        .position(|id| expected.iter().any(|e| e == *id))
        .map(|i| i + 1)
}

fn precision_at_k(window: &[&String], expected: &[String], k: usize) -> f64 {
    if k == 0 {
        return 0.0;
    }
    let intersect = window
        .iter()
        .filter(|id| expected.iter().any(|e| e == **id))
        .count();
    intersect as f64 / k as f64
}

fn recall_at_k(window: &[&String], expected: &[String]) -> f64 {
    if expected.is_empty() {
        return 0.0;
    }
    let intersect = window
        .iter()
        .filter(|id| expected.iter().any(|e| e == **id))
        .count();
    intersect as f64 / expected.len() as f64
}

fn ndcg_at_k(window: &[&String], expected: &[String], k: usize) -> f64 {
    if expected.is_empty() || k == 0 {
        return 0.0;
    }
    let dcg: f64 = window
        .iter()
        .enumerate()
        .map(|(idx, id)| {
            let relevance = if expected.iter().any(|e| e == *id) {
                1.0
            } else {
                0.0
            };
            relevance / ((idx + 2) as f64).log2()
        })
        .sum();
    let ideal_hits = expected.len().min(k);
    let idcg: f64 = (0..ideal_hits)
        .map(|idx| 1.0 / ((idx + 2) as f64).log2())
        .sum();
    if idcg == 0.0 {
        0.0
    } else {
        dcg / idcg
    }
}

fn mean(iter: impl IntoIterator<Item = f64>) -> f64 {
    let (sum, count) = iter
        .into_iter()
        .fold((0.0, 0usize), |(s, c), v| (s + v, c + 1));
    if count == 0 {
        0.0
    } else {
        sum / count as f64
    }
}

#[cfg(test)]
mod tests;