recallbench 0.2.0

A universal benchmark harness for AI memory systems
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};

/// Pipeline stages for checkpointing.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Stage {
    Ingest,
    Retrieve,
    Generate,
    Judge,
    Complete,
}

/// Checkpoint data saved after each pipeline stage.
#[derive(Debug, Serialize, Deserialize)]
pub struct Checkpoint {
    pub question_id: String,
    pub stage: Stage,
    pub ingest_ms: Option<u64>,
    pub retrieval_context: Option<String>,
    pub retrieval_ms: Option<u64>,
    pub hypothesis: Option<String>,
    pub generation_ms: Option<u64>,
    pub is_correct: Option<bool>,
    pub judge_ms: Option<u64>,
}

/// Get the checkpoint directory for a results file.
pub fn checkpoint_dir(results_path: &Path) -> PathBuf {
    let parent = results_path.parent().unwrap_or(Path::new("."));
    parent.join(".checkpoints")
}

/// Save a checkpoint for a question.
pub fn save_checkpoint(results_path: &Path, checkpoint: &Checkpoint) -> Result<()> {
    let dir = checkpoint_dir(results_path);
    std::fs::create_dir_all(&dir)?;

    let path = dir.join(format!("{}.json", checkpoint.question_id));
    let json = serde_json::to_string_pretty(checkpoint)?;
    std::fs::write(&path, json).context("Failed to write checkpoint")?;
    Ok(())
}

/// Load a checkpoint for a question, if one exists.
pub fn load_checkpoint(results_path: &Path, question_id: &str) -> Result<Option<Checkpoint>> {
    let path = checkpoint_dir(results_path).join(format!("{question_id}.json"));
    if !path.exists() {
        return Ok(None);
    }
    let content = std::fs::read_to_string(&path)?;
    let checkpoint: Checkpoint = serde_json::from_str(&content)?;
    Ok(Some(checkpoint))
}

/// Remove a checkpoint after question is fully evaluated.
pub fn remove_checkpoint(results_path: &Path, question_id: &str) -> Result<()> {
    let path = checkpoint_dir(results_path).join(format!("{question_id}.json"));
    if path.exists() {
        std::fs::remove_file(&path)?;
    }
    Ok(())
}

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

    #[test]
    fn checkpoint_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let results_path = dir.path().join("results.jsonl");

        let cp = Checkpoint {
            question_id: "q001".to_string(),
            stage: Stage::Retrieve,
            ingest_ms: Some(12),
            retrieval_context: Some("context here".to_string()),
            retrieval_ms: Some(8),
            hypothesis: None,
            generation_ms: None,
            is_correct: None,
            judge_ms: None,
        };

        save_checkpoint(&results_path, &cp).unwrap();
        let loaded = load_checkpoint(&results_path, "q001").unwrap().unwrap();
        assert_eq!(loaded.question_id, "q001");
        assert_eq!(loaded.stage, Stage::Retrieve);
        assert_eq!(loaded.retrieval_context, Some("context here".to_string()));

        remove_checkpoint(&results_path, "q001").unwrap();
        assert!(load_checkpoint(&results_path, "q001").unwrap().is_none());
    }

    #[test]
    fn load_nonexistent() {
        let dir = tempfile::tempdir().unwrap();
        let results_path = dir.path().join("results.jsonl");
        assert!(load_checkpoint(&results_path, "nonexistent").unwrap().is_none());
    }
}