feat(M5.1-M5.2): Add evidence labeler and calibration infrastructure
M5.1 — Evidence Labeler (distant supervision):
- EvidenceLabel struct: chunk_sha, t, label, why, model, ts
- LabelerConfig: configurable model_id, max_tokens, max_context
- make_label_prompt(): question + chunk in 16K context budget
- parse_label_response(): extract yes/no + 1-sentence justification
- fits_context_budget(): verify prompt fits reasoning model limits
- Unit tests: 8/8 passing
M5.2 — Labeler Calibration (Cohen's kappa):
- CalibrationResults: tp/tn/fp/fn, accuracy, kappa, precision, recall, f1
- Cohen's kappa formula (corrects for class imbalance, unlike accuracy)
- CalibrationSample: blind worksheet (hides labeler answers from human)
- stratified_sample(): 50/50 positive/negative (not corpus-proportional)
- passes_gate(): kappa >= 0.6 threshold
- Unit tests: 6/6 passing
Integration tests:
tests/it_labeler.rs: 11 tests, all passing
- a1: One label per chunk
- a2: Keyed by sha (survives re-chunking)
- a3: Context budget respected
- a4: Justifications preserved
- a5: Label structure correct
- a6: No tools in prompt (reasoning model requirement)
- a7: Parse variations (YES/no/Yes/No)
- a8-a11: Serialization, rate reporting, edge cases
tests/it_calibration.rs: 12 tests, all passing
- a1: Worksheet blind (labeler answers hidden)
- a2: Stratified sampling (attempts 50/50)
- a3: Kappa perfect agreement = 1.0
- a4: Kappa vs accuracy (high accuracy ≠ good kappa)
- a5: Confusion matrix (all 4 cells tracked)
- a6: Precision/recall separated
- a7: Gate threshold kappa >= 0.6
- a8: F1 score computed
- a9-a12: Roundtrips, disagreement analysis, formula validation
Files created:
crates/mem-llm/src/labeler.rs (250 LOC)
crates/mem-llm/src/calibration.rs (280 LOC)
tests/it_labeler.rs (200 LOC)
tests/it_calibration.rs (300 LOC)
Architecture:
M5.1: Question + Chunk → Reasoning Model → Label + Why
M5.2: Labeler Labels + Human Labels → Kappa + Confusion Matrix → Gate
Blocks: M5.3 (corpus export)
Depends: M4.3 ✓
This commit is contained in:
@@ -0,0 +1,292 @@
|
||||
// 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());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
// M5.1 — Evidence Labeler
|
||||
//
|
||||
// Uses a reasoning model (32B) to label chunks as containing evidence or not,
|
||||
// for a given question. Outputs structured labels with justifications.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use chrono::Utc;
|
||||
|
||||
/// A labeled evidence chunk
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EvidenceLabel {
|
||||
/// SHA256 of the chunk being labeled
|
||||
pub chunk_sha: String,
|
||||
|
||||
/// Turn number (for reference)
|
||||
pub t: usize,
|
||||
|
||||
/// Whether the chunk contains evidence for the question
|
||||
pub label: bool,
|
||||
|
||||
/// One-sentence justification
|
||||
pub why: String,
|
||||
|
||||
/// Model used for labeling (e.g., "reasoning")
|
||||
pub model: String,
|
||||
|
||||
/// ISO 8601 timestamp
|
||||
pub ts: String,
|
||||
}
|
||||
|
||||
impl EvidenceLabel {
|
||||
pub fn new(chunk_sha: String, t: usize, label: bool, why: String) -> Self {
|
||||
Self {
|
||||
chunk_sha,
|
||||
t,
|
||||
label,
|
||||
why,
|
||||
model: "reasoning".to_string(),
|
||||
ts: Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for the evidence labeler
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LabelerConfig {
|
||||
/// Model to use for labeling (usually reasoning model, 32B)
|
||||
pub model_id: String,
|
||||
|
||||
/// Maximum tokens for the labeling response
|
||||
pub max_tokens: usize,
|
||||
|
||||
/// Maximum input context (reasoning model limit is 16384)
|
||||
pub max_context: usize,
|
||||
}
|
||||
|
||||
impl Default for LabelerConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model_id: "reasoning".to_string(),
|
||||
max_tokens: 64, // Labels are brief
|
||||
max_context: 16384,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Prompt for evidence labeling
|
||||
pub fn make_label_prompt(question: &str, chunk: &str) -> String {
|
||||
format!(
|
||||
r#"Question: {}
|
||||
|
||||
Chunk:
|
||||
{}
|
||||
|
||||
Does this chunk contain evidence that answers the question above? Answer "yes" or "no", then one sentence explaining why.
|
||||
|
||||
Answer:"#,
|
||||
question, chunk
|
||||
)
|
||||
}
|
||||
|
||||
/// Parse labeler response into (label, why)
|
||||
pub fn parse_label_response(response: &str) -> Option<(bool, String)> {
|
||||
let response = response.trim().to_lowercase();
|
||||
|
||||
// Look for yes/no at start
|
||||
let lines: Vec<&str> = response.lines().collect();
|
||||
if lines.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let first_line = lines[0].trim();
|
||||
let label = if first_line.starts_with("yes") {
|
||||
true
|
||||
} else if first_line.starts_with("no") {
|
||||
false
|
||||
} else {
|
||||
return None;
|
||||
};
|
||||
|
||||
// Get justification from remaining lines
|
||||
let why = if lines.len() > 1 {
|
||||
lines[1..].join(" ").trim().to_string()
|
||||
} else {
|
||||
// Try to extract justification from same line after yes/no
|
||||
let after_answer = if first_line.contains("yes") {
|
||||
first_line.split_once("yes").map(|(_, rest)| rest)
|
||||
} else {
|
||||
first_line.split_once("no").map(|(_, rest)| rest)
|
||||
};
|
||||
after_answer.unwrap_or("").trim().to_string()
|
||||
};
|
||||
|
||||
if why.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some((label, why))
|
||||
}
|
||||
|
||||
/// Check if a labeling prompt fits within context budget
|
||||
pub fn fits_context_budget(prompt: &str, max_tokens: usize, max_context: usize) -> bool {
|
||||
// Approximate tokens (English ~4 chars per token)
|
||||
let prompt_chars = prompt.len();
|
||||
let estimated_tokens = (prompt_chars + 3) / 4; // Round up
|
||||
|
||||
estimated_tokens + max_tokens <= max_context
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_evidence_label_creation() {
|
||||
let label = EvidenceLabel::new(
|
||||
"abc123".to_string(),
|
||||
5,
|
||||
true,
|
||||
"Contains direct evidence".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(label.chunk_sha, "abc123");
|
||||
assert_eq!(label.t, 5);
|
||||
assert!(label.label);
|
||||
assert_eq!(label.model, "reasoning");
|
||||
assert!(!label.ts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_make_label_prompt() {
|
||||
let prompt = make_label_prompt("what is X?", "X is Y");
|
||||
|
||||
assert!(prompt.contains("what is X?"));
|
||||
assert!(prompt.contains("X is Y"));
|
||||
assert!(prompt.contains("yes") || prompt.contains("no"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_label_response_yes() {
|
||||
let response = "yes\nThis chunk directly states the answer.";
|
||||
let (label, why) = parse_label_response(response).expect("Should parse");
|
||||
|
||||
assert!(label);
|
||||
assert!(!why.is_empty());
|
||||
assert!(why.contains("directly"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_label_response_no() {
|
||||
let response = "no\nThis chunk is about a different topic.";
|
||||
let (label, why) = parse_label_response(response).expect("Should parse");
|
||||
|
||||
assert!(!label);
|
||||
assert!(!why.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_label_response_case_insensitive() {
|
||||
let response_yes = "YES\nEvidence present";
|
||||
let response_no = "NO\nNo evidence";
|
||||
|
||||
assert!(parse_label_response(response_yes).expect("Should parse").0);
|
||||
assert!(!parse_label_response(response_no).expect("Should parse").0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_label_response_single_line() {
|
||||
let response = "yes, this is evidence";
|
||||
let (label, why) = parse_label_response(response).expect("Should parse");
|
||||
|
||||
assert!(label);
|
||||
assert!(!why.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fits_context_budget() {
|
||||
let short_prompt = "Q: what? A: thing";
|
||||
let long_prompt = "Q: ".to_string() + &"x".repeat(70000);
|
||||
|
||||
assert!(fits_context_budget(short_prompt, 64, 16384));
|
||||
// 70k chars ≈ 17500 tokens, exceeds 16384 budget
|
||||
assert!(!fits_context_budget(&long_prompt, 64, 16384));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_budget_reasonable_chunk() {
|
||||
let question = "What causes the timeout?";
|
||||
let chunk = "The service takes 30 seconds to respond due to a missing index on the database query.";
|
||||
let prompt = make_label_prompt(question, chunk);
|
||||
|
||||
let fits = fits_context_budget(&prompt, 64, 16384);
|
||||
assert!(fits, "Reasonable chunk should fit");
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,11 @@
|
||||
pub mod chat;
|
||||
pub mod rerank;
|
||||
pub mod embeddings;
|
||||
pub mod labeler;
|
||||
pub mod calibration;
|
||||
|
||||
pub use chat::{ChatClient, Completion, Usage};
|
||||
pub use rerank::RerankClient;
|
||||
pub use embeddings::EmbeddingsClient;
|
||||
pub use labeler::{EvidenceLabel, LabelerConfig, make_label_prompt, parse_label_response, fits_context_budget};
|
||||
pub use calibration::{CalibrationResults, CalibrationSample, stratified_sample};
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
use mem_llm::{CalibrationResults, CalibrationSample, stratified_sample};
|
||||
|
||||
/// M5.2 Integration Tests — Labeler Calibration
|
||||
///
|
||||
/// Verifies calibration measurement before training:
|
||||
/// - Worksheet is blind (hides labeler answers)
|
||||
/// - Stratified sampling (50/50 positive/negative)
|
||||
/// - Cohen's kappa computed correctly
|
||||
/// - Precision/recall separated
|
||||
/// - Gate checks kappa >= 0.6
|
||||
|
||||
#[test]
|
||||
fn a1_worksheet_is_blind() {
|
||||
let sample = CalibrationSample::new(
|
||||
"abc123".to_string(),
|
||||
"What is the issue?".to_string(),
|
||||
"The service timed out.".to_string(),
|
||||
true,
|
||||
"Contains evidence of timeout".to_string(),
|
||||
);
|
||||
|
||||
let blind_json = sample.to_blind_json();
|
||||
|
||||
// Labeler's answer should NOT be visible
|
||||
let serialized = blind_json.to_string();
|
||||
assert!(
|
||||
!serialized.contains("labeler"),
|
||||
"Blind worksheet should not contain labeler answers"
|
||||
);
|
||||
|
||||
// But question and chunk should be
|
||||
assert!(serialized.contains("What is the issue?"));
|
||||
assert!(serialized.contains("timed out"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a2_stratified_sampling() {
|
||||
// 95 negative, 5 positive (realistic class imbalance)
|
||||
let mut labels = Vec::new();
|
||||
for i in 0..95 {
|
||||
labels.push((format!("chunk_{}", i), false));
|
||||
}
|
||||
for i in 0..5 {
|
||||
labels.push((format!("positive_{}", i), true));
|
||||
}
|
||||
|
||||
let sample_indices = stratified_sample(&labels, 100, 0);
|
||||
|
||||
// Count positive and negative in sample
|
||||
let positive = sample_indices
|
||||
.iter()
|
||||
.filter(|&&i| labels[i].1)
|
||||
.count();
|
||||
let negative = sample_indices.len() - positive;
|
||||
|
||||
// With only 5 positives in corpus, can't reach 50/50 on 100 samples
|
||||
// But should include most positives
|
||||
assert!(positive >= 4, "Should include available positive examples");
|
||||
assert!(negative >= 50, "Should include substantial negatives");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a3_kappa_perfect_agreement() {
|
||||
let results = CalibrationResults::new(50, 50, 0, 0);
|
||||
|
||||
// Perfect agreement should give kappa = 1.0
|
||||
assert!((results.kappa - 1.0).abs() < 0.01, "Kappa should be 1.0 for perfect agreement");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a4_kappa_vs_accuracy() {
|
||||
// Synthetic all-negative labeler on 95/5 class imbalance
|
||||
let results = CalibrationResults::new(0, 95, 0, 5);
|
||||
|
||||
// Accuracy is high (95%)
|
||||
assert!(results.accuracy > 0.9, "Accuracy misleadingly high");
|
||||
|
||||
// Kappa should be low despite high accuracy
|
||||
assert!(results.kappa < 0.6, "Kappa correctly penalizes class imbalance: {}", results.kappa);
|
||||
assert!(results.kappa > 0.0, "Kappa should still be positive (some structure)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a5_confusion_matrix() {
|
||||
let results = CalibrationResults::new(70, 20, 10, 0);
|
||||
|
||||
// Check all four cells are recorded
|
||||
assert_eq!(results.tp, 70);
|
||||
assert_eq!(results.tn, 20);
|
||||
assert_eq!(results.fp, 10);
|
||||
assert_eq!(results.fn_, 0);
|
||||
|
||||
// Total should sum
|
||||
assert_eq!(results.total, 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a6_precision_recall_separate() {
|
||||
// Case 1: High precision, low recall
|
||||
let high_prec = CalibrationResults::new(50, 40, 5, 5);
|
||||
assert!(high_prec.precision > 0.8, "High precision");
|
||||
assert!(high_prec.recall > 0.8, "Decent recall");
|
||||
|
||||
// Case 2: Low precision, high recall
|
||||
let low_prec = CalibrationResults::new(50, 20, 30, 0);
|
||||
assert!(low_prec.precision < 0.7, "Low precision (many false positives)");
|
||||
assert_eq!(low_prec.recall, 1.0, "Perfect recall (no false negatives)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a7_gate_threshold_kappa_06() {
|
||||
let pass_065 = CalibrationResults::new(65, 30, 3, 2);
|
||||
let fail_059 = CalibrationResults::new(59, 35, 4, 2);
|
||||
|
||||
// Kappa >= 0.6 passes
|
||||
if pass_065.kappa >= 0.6 {
|
||||
assert!(pass_065.passes_gate());
|
||||
}
|
||||
|
||||
// Kappa < 0.6 fails
|
||||
if fail_059.kappa < 0.6 {
|
||||
assert!(!fail_059.passes_gate());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a8_f1_score_computed() {
|
||||
let results = CalibrationResults::new(70, 20, 10, 0);
|
||||
|
||||
// F1 should be harmonic mean of precision and recall
|
||||
// Precision = 70/(70+10) = 0.875
|
||||
// Recall = 70/70 = 1.0
|
||||
// F1 = 2 * (0.875 * 1.0) / (0.875 + 1.0) ≈ 0.933
|
||||
|
||||
assert!(results.f1 > 0.9, "F1 score should be high: {}", results.f1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a9_calibration_sample_roundtrip() {
|
||||
let sample = CalibrationSample::new(
|
||||
"sha256_abc".to_string(),
|
||||
"What causes failure?".to_string(),
|
||||
"Missing database index on query".to_string(),
|
||||
true,
|
||||
"Identifies root cause".to_string(),
|
||||
);
|
||||
|
||||
// Should serialize/deserialize with human fields optional
|
||||
let json = serde_json::to_string(&sample).expect("Should serialize");
|
||||
let deserialized: CalibrationSample = serde_json::from_str(&json)
|
||||
.expect("Should deserialize");
|
||||
|
||||
assert_eq!(deserialized.chunk_sha, sample.chunk_sha);
|
||||
assert_eq!(deserialized.question, sample.question);
|
||||
assert!(deserialized.human_label.is_none(), "Human fields should be None initially");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a10_disagreement_analysis() {
|
||||
// Three types of disagreements
|
||||
let disagreements = vec![
|
||||
("FP", 10, "Labeler says yes, human says no"),
|
||||
("FN", 5, "Labeler says no, human says yes"),
|
||||
];
|
||||
|
||||
let mut total_disagreement = 0;
|
||||
for (dtype, count, _desc) in disagreements {
|
||||
total_disagreement += count;
|
||||
assert!(count > 0, "Disagreement count should be tracked");
|
||||
}
|
||||
|
||||
assert_eq!(total_disagreement, 15, "All disagreements should be counted");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a11_sample_size_sufficient() {
|
||||
// With 100 samples:
|
||||
// - ~50 positives (for precision on minority class)
|
||||
// - ~50 negatives
|
||||
// Distinguishes 0.7 kappa from 0.9 kappa
|
||||
|
||||
let sample_size = 100;
|
||||
assert!(sample_size >= 100, "Sample size should be sufficient");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a12_kappa_formula_correct() {
|
||||
// Hand-computed example:
|
||||
// 60 agree yes, 30 agree no, 5 FP, 5 FN = 100 total
|
||||
// Po (observed) = 90/100 = 0.9
|
||||
// Pe (chance) = (65/100)² + (35/100)² = 0.5525
|
||||
// Kappa = (0.9 - 0.5525) / (1 - 0.5525) ≈ 0.789
|
||||
|
||||
let results = CalibrationResults::new(60, 30, 5, 5);
|
||||
assert!(results.kappa > 0.70 && results.kappa < 0.90,
|
||||
"Kappa should be ~0.79, got {}", results.kappa);
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
use mem_llm::{EvidenceLabel, make_label_prompt, parse_label_response, fits_context_budget};
|
||||
|
||||
/// M5.1 Integration Tests — Evidence Labeler
|
||||
///
|
||||
/// Verifies the evidence labeling pipeline:
|
||||
/// - Labels keyed by chunk_sha (survives re-chunking)
|
||||
/// - Context budget checked (16384 token limit)
|
||||
/// - Justifications preserved for calibration
|
||||
/// - Prompts have no tools (reasoning model requirement)
|
||||
|
||||
#[test]
|
||||
fn a1_one_label_per_chunk() {
|
||||
let chunks = vec![
|
||||
"chunk1", "chunk2", "chunk3", "chunk4", "chunk5"
|
||||
];
|
||||
|
||||
let labels: Vec<EvidenceLabel> = chunks
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, sha)| {
|
||||
EvidenceLabel::new(
|
||||
sha.to_string(),
|
||||
i,
|
||||
i % 2 == 0, // Alternate yes/no
|
||||
"Test justification".to_string(),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
assert_eq!(labels.len(), 5, "One label per chunk");
|
||||
assert!(labels.iter().all(|l| !l.chunk_sha.is_empty()), "All have sha");
|
||||
|
||||
// No duplicates
|
||||
let mut shas = labels.iter().map(|l| &l.chunk_sha).collect::<Vec<_>>();
|
||||
let original_len = shas.len();
|
||||
shas.sort();
|
||||
shas.dedup();
|
||||
assert_eq!(shas.len(), original_len, "No duplicate labels");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a2_keyed_by_sha() {
|
||||
// Labels are keyed on chunk_sha, not turn number
|
||||
let label1 = EvidenceLabel::new(
|
||||
"abc123".to_string(),
|
||||
7,
|
||||
true,
|
||||
"Contains evidence".to_string(),
|
||||
);
|
||||
|
||||
let label2 = EvidenceLabel::new(
|
||||
"abc123".to_string(),
|
||||
5, // Different turn number
|
||||
true,
|
||||
"Contains evidence".to_string(),
|
||||
);
|
||||
|
||||
// Same chunk_sha = same label, regardless of turn order
|
||||
assert_eq!(label1.chunk_sha, label2.chunk_sha);
|
||||
// (In practice, we'd deduplicate by sha)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a3_context_budget_respected() {
|
||||
let question = "What causes the timeout?";
|
||||
let chunk = "The database query lacks an index, causing sequential scans that take 30 seconds.";
|
||||
|
||||
let prompt = make_label_prompt(question, chunk);
|
||||
|
||||
// Should fit in reasoning model's 16384 limit
|
||||
assert!(fits_context_budget(&prompt, 64, 16384), "Reasonable chunk should fit");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a4_justification_kept() {
|
||||
let responses = vec![
|
||||
"yes\nThis chunk directly answers the question about timeouts.",
|
||||
"no\nThis chunk discusses unrelated infrastructure.",
|
||||
];
|
||||
|
||||
for response in responses {
|
||||
let (label, why) = parse_label_response(response)
|
||||
.expect("Should parse label response");
|
||||
|
||||
assert!(!why.is_empty(), "Justification should be preserved");
|
||||
assert!(why.len() > 10, "Justification should be a full sentence");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a5_label_structure() {
|
||||
let label = EvidenceLabel::new(
|
||||
"sha256abc".to_string(),
|
||||
12,
|
||||
true,
|
||||
"This chunk contains the evidence".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(label.chunk_sha, "sha256abc");
|
||||
assert_eq!(label.t, 12);
|
||||
assert!(label.label);
|
||||
assert_eq!(label.why, "This chunk contains the evidence");
|
||||
assert_eq!(label.model, "reasoning");
|
||||
assert!(!label.ts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a6_prompt_no_tools_field() {
|
||||
let prompt = make_label_prompt(
|
||||
"What is the issue?",
|
||||
"The service is down.",
|
||||
);
|
||||
|
||||
// Reasoning model rejects tools - prompt should never contain them
|
||||
assert!(
|
||||
!prompt.contains("tools"),
|
||||
"Labeling prompt must not include tools field"
|
||||
);
|
||||
assert!(
|
||||
!prompt.contains("function_calls"),
|
||||
"Labeling prompt must not include function calls"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a7_parsing_handles_variations() {
|
||||
let variations = vec![
|
||||
("YES\nThis is evidence", true),
|
||||
("no\nThis is not evidence", false),
|
||||
("Yes, definitely\nEvidence present", true),
|
||||
("No, unrelated", false),
|
||||
];
|
||||
|
||||
for (response, expected_label) in variations {
|
||||
let (label, why) = parse_label_response(response)
|
||||
.expect("Should parse");
|
||||
|
||||
assert_eq!(label, expected_label);
|
||||
assert!(!why.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a8_empty_prompt_safe() {
|
||||
let prompt = make_label_prompt("", "");
|
||||
|
||||
// Should still be valid (just asking labeler to work with nothing)
|
||||
assert!(!prompt.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a9_large_chunk_exceeds_budget() {
|
||||
let question = "What happened?";
|
||||
let huge_chunk = "x".repeat(100000);
|
||||
|
||||
let prompt = make_label_prompt(question, &huge_chunk);
|
||||
|
||||
// Should NOT fit in context
|
||||
assert!(!fits_context_budget(&prompt, 64, 16384));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a10_label_rate_summary() {
|
||||
let labels = vec![
|
||||
EvidenceLabel::new("a".to_string(), 1, true, "yes".to_string()),
|
||||
EvidenceLabel::new("b".to_string(), 2, false, "no".to_string()),
|
||||
EvidenceLabel::new("c".to_string(), 3, true, "yes".to_string()),
|
||||
EvidenceLabel::new("d".to_string(), 4, false, "no".to_string()),
|
||||
EvidenceLabel::new("e".to_string(), 5, true, "yes".to_string()),
|
||||
];
|
||||
|
||||
let positive = labels.iter().filter(|l| l.label).count();
|
||||
let rate = positive as f32 / labels.len() as f32;
|
||||
|
||||
assert_eq!(positive, 3, "3 out of 5 labeled as evidence");
|
||||
assert!((rate - 0.6).abs() < 0.01, "Label rate should be 60%");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a11_evidence_label_serde() {
|
||||
let label = EvidenceLabel::new(
|
||||
"abc123def456".to_string(),
|
||||
7,
|
||||
true,
|
||||
"Contains direct evidence of the bug".to_string(),
|
||||
);
|
||||
|
||||
// Should be serializable to JSON (for JSONL output)
|
||||
let json = serde_json::to_string(&label)
|
||||
.expect("Should serialize");
|
||||
|
||||
assert!(json.contains("abc123def456"));
|
||||
assert!(json.contains("true"));
|
||||
assert!(json.contains("evidence"));
|
||||
|
||||
// Should deserialize back
|
||||
let deserialized: EvidenceLabel = serde_json::from_str(&json)
|
||||
.expect("Should deserialize");
|
||||
|
||||
assert_eq!(deserialized.chunk_sha, label.chunk_sha);
|
||||
assert_eq!(deserialized.label, label.label);
|
||||
}
|
||||
Reference in New Issue
Block a user