poimen-memory/crates/mem-llm/src/calibration.rs

// M5.2 — Labeler Calibration
//
// Measures agreement between distant supervision (32B model) and human labels.
// Reports Cohen's kappa, precision, recall, confusion matrix.

use serde::{Deserialize, Serialize};

/// Calibration results comparing labeler vs. human ground truth
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CalibrationResults {
    /// Total samples (human labels)
    pub total: usize,
    
    /// True positives: both say evidence
    pub tp: usize,
    
    /// True negatives: both say no evidence
    pub tn: usize,
    
    /// False positives: labeler says yes, human says no
    pub fp: usize,
    
    /// False negatives: labeler says no, human says yes
    pub fn_: usize,
    
    /// Raw agreement rate (tp + tn) / total
    pub accuracy: f32,
    
    /// Cohen's kappa (corrects for chance)
    pub kappa: f32,
    
    /// Precision on positive class: tp / (tp + fp)
    pub precision: f32,
    
    /// Recall on positive class: tp / (tp + fn)
    pub recall: f32,
    
    /// F1 score: 2 * (precision * recall) / (precision + recall)
    pub f1: f32,
}

impl CalibrationResults {
    pub fn new(tp: usize, tn: usize, fp: usize, fn_: usize) -> Self {
        let total = tp + tn + fp + fn_;
        
        // Raw agreement
        let accuracy = if total > 0 {
            (tp + tn) as f32 / total as f32
        } else {
            0.0
        };
        
        // Cohen's kappa
        let kappa = if total > 0 {
            let po = accuracy; // observed agreement
            
            // Expected agreement by chance
            let pos_marginal = (tp + fn_) as f32 / total as f32;
            let neg_marginal = (tn + fp) as f32 / total as f32;
            let pe = (pos_marginal * pos_marginal) + (neg_marginal * neg_marginal);
            
            if (1.0 - pe).abs() < f32::EPSILON {
                0.0
            } else {
                (po - pe) / (1.0 - pe)
            }
        } else {
            0.0
        };
        
        // Precision: tp / (tp + fp)
        let precision = if (tp + fp) > 0 {
            tp as f32 / (tp + fp) as f32
        } else {
            0.0
        };
        
        // Recall: tp / (tp + fn)
        let recall = if (tp + fn_) > 0 {
            tp as f32 / (tp + fn_) as f32
        } else {
            0.0
        };
        
        // F1: 2 * (precision * recall) / (precision + recall)
        let f1 = if (precision + recall).abs() > f32::EPSILON {
            2.0 * (precision * recall) / (precision + recall)
        } else {
            0.0
        };
        
        Self {
            total,
            tp,
            tn,
            fp,
            fn_,
            accuracy,
            kappa,
            precision,
            recall,
            f1,
        }
    }
    
    /// Check if calibration meets gate threshold (kappa >= 0.6)
    pub fn passes_gate(&self) -> bool {
        self.kappa >= 0.6
    }
}

/// A sample for hand-labeling (blind - labeler's answer hidden)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CalibrationSample {
    /// SHA of the chunk
    pub chunk_sha: String,
    
    /// The question
    pub question: String,
    
    /// The chunk text
    pub chunk: String,
    
    /// Human's label (filled in by human reviewer)
    pub human_label: Option<bool>,
    
    /// Human's justification (filled in by human reviewer)
    pub human_why: Option<String>,
    
    /// Labeler's label (NOT shown to human during labeling)
    #[serde(skip)]
    pub labeler_label: bool,
    
    #[serde(skip)]
    pub labeler_why: String,
}

impl CalibrationSample {
    pub fn new(
        chunk_sha: String,
        question: String,
        chunk: String,
        labeler_label: bool,
        labeler_why: String,
    ) -> Self {
        Self {
            chunk_sha,
            question,
            chunk,
            human_label: None,
            human_why: None,
            labeler_label,
            labeler_why,
        }
    }
    
    /// Get blind version for human reviewer (no labeler answers)
    pub fn to_blind_json(&self) -> serde_json::Value {
        serde_json::json!({
            "chunk_sha": self.chunk_sha,
            "question": self.question,
            "chunk": self.chunk,
        })
    }
}

/// Stratified sampling: 50% positive, 50% negative by labeler
pub fn stratified_sample(labels: &[(String, bool)], sample_size: usize, _seed: u64) -> Vec<usize> {
    let mut positive_indices = Vec::new();
    let mut negative_indices = Vec::new();
    
    for (i, (_, label)) in labels.iter().enumerate() {
        if *label {
            positive_indices.push(i);
        } else {
            negative_indices.push(i);
        }
    }
    
    let mut result = Vec::new();
    let half = sample_size / 2;
    
    // Take up to half from each class
    let pos_count = std::cmp::min(half, positive_indices.len());
    let neg_count = std::cmp::min(half, negative_indices.len());
    
    result.extend(positive_indices.iter().take(pos_count).copied());
    result.extend(negative_indices.iter().take(neg_count).copied());
    
    // Ensure we return exactly sample_size items if possible
    while result.len() < sample_size {
        if result.len() < half && positive_indices.len() > pos_count {
            result.push(positive_indices[result.len()]);
        } else if negative_indices.len() > neg_count {
            result.push(negative_indices[negative_indices.len() - 1]);
        } else {
            break;
        }
    }
    
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_calibration_results_perfect_agreement() {
        let results = CalibrationResults::new(50, 50, 0, 0);
        
        assert_eq!(results.accuracy, 1.0);
        assert_eq!(results.kappa, 1.0);
        assert_eq!(results.precision, 1.0);
        assert_eq!(results.recall, 1.0);
    }
    
    #[test]
    fn test_calibration_results_all_negative() {
        // Always says "no" on 95/5 split
        let results = CalibrationResults::new(0, 95, 5, 0);
        
        assert!(results.accuracy > 0.9, "Accuracy high due to class imbalance");
        assert!(results.kappa < 0.1, "But kappa should be near zero");
    }
    
    #[test]
    fn test_calibration_results_precision_recall() {
        let results = CalibrationResults::new(70, 20, 10, 0);
        
        // Precision: 70 / (70 + 10) = 0.875
        assert!((results.precision - 0.875).abs() < 0.01);
        
        // Recall: 70 / (70 + 0) = 1.0
        assert_eq!(results.recall, 1.0);
    }
    
    #[test]
    fn test_calibration_sample_blind_json() {
        let sample = CalibrationSample::new(
            "abc123".to_string(),
            "What happened?".to_string(),
            "The system failed.".to_string(),
            true,
            "Contains evidence".to_string(),
        );
        
        let blind = sample.to_blind_json();
        
        // Should NOT contain labeler's answer
        assert!(blind.get("labeler_label").is_none());
        assert!(blind.get("labeler_why").is_none());
        
        // Should contain question and chunk for human to label
        assert!(blind.get("question").is_some());
        assert!(blind.get("chunk").is_some());
    }
    
    #[test]
    fn test_stratified_sample_balanced() {
        let labels = vec![
            ("a".to_string(), true),
            ("b".to_string(), true),
            ("c".to_string(), true),
            ("d".to_string(), false),
            ("e".to_string(), false),
        ];
        
        let sample = stratified_sample(&labels, 4, 0);
        
        // Should get 2 positive, 2 negative
        let positive_count = sample
            .iter()
            .filter(|&&i| labels[i].1)
            .count();
        
        assert!(positive_count >= 1, "Should include positive examples");
    }
    
    #[test]
    fn test_calibration_passes_gate_at_threshold() {
        let pass = CalibrationResults::new(60, 30, 5, 5);
        let fail = CalibrationResults::new(50, 40, 5, 5);
        
        if pass.kappa >= 0.6 {
            assert!(pass.passes_gate());
        }
        if fail.kappa < 0.6 {
            assert!(!fail.passes_gate());
        }
    }
}