klieo_eval/

lib.rs

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
//! 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 std::collections::HashSet;

use serde::Serialize;

#[cfg(feature = "agent-eval")]
pub mod agent_eval;
#[cfg(feature = "agent-eval")]
#[allow(deprecated)]
pub use agent_eval::eval_capture;
#[cfg(feature = "agent-eval")]
pub use agent_eval::{check_regression, eval_capture_determinism, EvalMetrics, Regression};
#[cfg(feature = "agent-eval")]
pub mod live_eval;
#[cfg(feature = "agent-eval")]
pub use live_eval::{eval_capture_live, LiveEvalMetrics};

pub mod classifier_eval;
pub use classifier_eval::{
    extract_label, score_classification, ClassificationCase, ClassificationReport, LabelSpec,
    LabelStat, Miss,
};

/// 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 {
    /// Query text (used only for trace output; not consumed by the metrics).
    pub query: String,
    /// Ground-truth ids that should land in the top-K window.
    pub expected: Vec<String>,
    /// Actual ranked id list returned by the recall pipeline (top result first).
    pub returned_ids: Vec<String>,
}

impl RecallSample {
    /// Construct a sample. An empty `expected` makes recall and NDCG
    /// undefined — filter such samples before scoring with [`score_recall`].
    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 {
    /// Top-K window used for every per-query metric.
    pub k: usize,
    /// One [`PerQueryScore`] per input sample, in input order.
    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 {
    /// Query text, copied from the input sample for trace output.
    pub query: String,
    /// `|expected|` for this sample — used by aggregators to skip
    /// queries where recall / NDCG would be undefined.
    pub expected_count: usize,
    /// 1-based rank of the first expected id within the top-K window,
    /// or `None` if no expected id landed in the window.
    pub first_hit_rank: Option<usize>,
    /// `|expected ∩ top-K| / k` for this query (always defined; `0.0` when `k == 0`).
    pub precision_at_k: f64,
    /// `|expected ∩ top-K| / |expected|` for this query (`0.0` when `expected` is empty).
    pub recall_at_k: f64,
    /// Normalised DCG@K with binary relevance for this query
    /// (`0.0` when `expected` is empty or `k == 0`).
    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 expected_set: HashSet<&str> = sample.expected.iter().map(String::as_str).collect();
    let window: Vec<&str> = sample
        .returned_ids
        .iter()
        .take(k)
        .map(String::as_str)
        .collect();
    let expected_count = sample.expected.len();
    let first_hit_rank = first_hit_rank(&window, &expected_set);
    let precision_at_k = precision_at_k(&window, &expected_set, k);
    let recall_at_k = recall_at_k(&window, &expected_set, expected_count);
    let ndcg_at_k = ndcg_at_k(&window, &expected_set, expected_count, k);
    PerQueryScore {
        query: sample.query.clone(),
        expected_count,
        first_hit_rank,
        precision_at_k,
        recall_at_k,
        ndcg_at_k,
    }
}

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

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

fn recall_at_k(window: &[&str], expected: &HashSet<&str>, expected_count: usize) -> f64 {
    if expected_count == 0 {
        return 0.0;
    }
    let intersect = window.iter().filter(|id| expected.contains(**id)).count();
    intersect as f64 / expected_count as f64
}

fn ndcg_at_k(window: &[&str], expected: &HashSet<&str>, expected_count: usize, k: usize) -> f64 {
    if expected_count == 0 || k == 0 {
        return 0.0;
    }
    let dcg: f64 = window
        .iter()
        .enumerate()
        .map(|(idx, id)| {
            let relevance = if expected.contains(*id) { 1.0 } else { 0.0 };
            relevance / ((idx + 2) as f64).log2()
        })
        .sum();
    let ideal_hits = expected_count.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;