recallbench 0.1.0

A universal benchmark harness for AI memory systems
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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

pub mod convomem;
pub mod custom;
pub mod download;
pub mod halumem;
pub mod locomo;
pub mod longmemeval;
pub mod mab;
pub mod membench;

use std::collections::HashMap;

use anyhow::Result;

use crate::traits::BenchmarkDataset;

/// Registry mapping dataset names to loader functions.
pub struct DatasetRegistry {
    datasets: HashMap<String, DatasetInfo>,
}

/// Metadata about an available dataset.
#[derive(Debug, Clone)]
pub struct DatasetInfo {
    pub name: String,
    pub description: String,
    pub variants: Vec<String>,
}

impl DatasetRegistry {
    pub fn new() -> Self {
        let mut registry = Self {
            datasets: HashMap::new(),
        };
        registry.register_defaults();
        registry
    }

    fn register_defaults(&mut self) {
        self.datasets.insert("longmemeval".to_string(), DatasetInfo {
            name: "longmemeval".to_string(),
            description: "LongMemEval (ICLR 2025) — 500 questions, 5 memory abilities, 7 types".to_string(),
            variants: vec!["oracle".to_string(), "small".to_string(), "medium".to_string()],
        });
        self.datasets.insert("locomo".to_string(), DatasetInfo {
            name: "locomo".to_string(),
            description: "LoCoMo (Snap Research) — long-context conversation memory".to_string(),
            variants: vec!["default".to_string()],
        });
        self.datasets.insert("convomem".to_string(), DatasetInfo {
            name: "convomem".to_string(),
            description: "ConvoMem — conversational memory evaluation".to_string(),
            variants: vec!["default".to_string()],
        });
        self.datasets.insert("membench".to_string(), DatasetInfo {
            name: "membench".to_string(),
            description: "MemBench (ACL 2025) — multi-aspect memory evaluation".to_string(),
            variants: vec!["default".to_string()],
        });
        self.datasets.insert("memoryagentbench".to_string(), DatasetInfo {
            name: "memoryagentbench".to_string(),
            description: "MemoryAgentBench (ICLR 2026) — selective forgetting, fact consolidation".to_string(),
            variants: vec!["default".to_string()],
        });
        self.datasets.insert("halumem".to_string(), DatasetInfo {
            name: "halumem".to_string(),
            description: "HaluMem (MemTensor) — memory hallucination detection".to_string(),
            variants: vec!["medium".to_string(), "long".to_string()],
        });
    }

    /// List all available datasets.
    pub fn list(&self) -> Vec<&DatasetInfo> {
        let mut datasets: Vec<_> = self.datasets.values().collect();
        datasets.sort_by(|a, b| a.name.cmp(&b.name));
        datasets
    }

    /// Get info for a specific dataset.
    pub fn get(&self, name: &str) -> Option<&DatasetInfo> {
        self.datasets.get(name)
    }

    /// Load a dataset by name and variant.
    pub async fn load(&self, name: &str, variant: &str, force_download: bool) -> Result<Box<dyn BenchmarkDataset>> {
        match name {
            "longmemeval" => {
                let dataset = longmemeval::LongMemEvalDataset::load(variant, force_download).await?;
                Ok(Box::new(dataset))
            }
            // These datasets require the user to provide the data file path
            // since they don't have standardized download URLs
            _ => anyhow::bail!(
                "Dataset '{name}' requires a local file. Use 'recallbench validate <path>' with a custom dataset file, \
                or download the dataset manually and use --dataset custom."
            ),
        }
    }
}

impl Default for DatasetRegistry {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn registry_has_all_datasets() {
        let registry = DatasetRegistry::new();
        assert!(registry.get("longmemeval").is_some());
        assert!(registry.get("locomo").is_some());
        assert!(registry.get("convomem").is_some());
        assert!(registry.get("membench").is_some());
        assert!(registry.get("memoryagentbench").is_some());
        assert!(registry.get("halumem").is_some());
    }

    #[test]
    fn registry_list_sorted() {
        let registry = DatasetRegistry::new();
        let datasets = registry.list();
        assert_eq!(datasets.len(), 6);
        // Verify sorted
        let names: Vec<_> = datasets.iter().map(|d| d.name.as_str()).collect();
        let mut sorted = names.clone();
        sorted.sort();
        assert_eq!(names, sorted);
    }

    #[test]
    fn registry_unknown_returns_none() {
        let registry = DatasetRegistry::new();
        assert!(registry.get("nonexistent").is_none());
    }
}