lean_ctx/core/locomo/
dataset.rs

1//! LoCoMo benchmark dataset schema + loader (#291).
2//!
3//! Accepts both NDJSON (one [`LocomoSample`] per line, `#` comments allowed) and a
4//! plain JSON array, so the bundled reference suite and the official LoCoMo dataset
5//! can be loaded by the same code.
6
7use std::path::Path;
8
9use serde::{Deserialize, Serialize};
10
11/// One LoCoMo sample: a multi-session conversation plus its question/answer set.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct LocomoSample {
14    pub id: String,
15    pub sessions: Vec<Session>,
16    pub qa: Vec<QaItem>,
17}
18
19/// An ordered conversation session.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct Session {
22    #[serde(default)]
23    pub session_id: String,
24    pub turns: Vec<Turn>,
25}
26
27/// A single dialog turn.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct Turn {
30    pub speaker: String,
31    pub text: String,
32}
33
34/// A question with its acceptable gold answers and LoCoMo category
35/// (1=single-hop, 2=multi-hop, 3=temporal, 4=open-domain, 5=adversarial).
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct QaItem {
38    pub question: String,
39    pub answers: Vec<String>,
40    #[serde(default = "default_category")]
41    pub category: u8,
42}
43
44fn default_category() -> u8 {
45    1
46}
47
48impl LocomoSample {
49    /// Flattened transcript (`Speaker: text` per turn) — the baseline an agent
50    /// would otherwise dump into context wholesale.
51    pub fn transcript(&self) -> String {
52        let mut lines = Vec::new();
53        for session in &self.sessions {
54            for turn in &session.turns {
55                lines.push(format!("{}: {}", turn.speaker, turn.text));
56            }
57        }
58        lines.join("\n")
59    }
60
61    /// Total number of conversation turns across all sessions.
62    pub fn turn_count(&self) -> usize {
63        self.sessions.iter().map(|s| s.turns.len()).sum()
64    }
65}
66
67/// Parse a suite from raw text (NDJSON or JSON array).
68pub fn parse_suite(raw: &str) -> Result<Vec<LocomoSample>, String> {
69    let trimmed = raw.trim_start();
70    if trimmed.starts_with('[') {
71        return serde_json::from_str(trimmed).map_err(|e| format!("invalid JSON array: {e}"));
72    }
73    let mut out = Vec::new();
74    for (i, line) in raw.lines().enumerate() {
75        let l = line.trim();
76        if l.is_empty() || l.starts_with('#') {
77            continue;
78        }
79        let sample: LocomoSample =
80            serde_json::from_str(l).map_err(|e| format!("line {}: {e}", i + 1))?;
81        out.push(sample);
82    }
83    if out.is_empty() {
84        return Err("suite contained no samples".to_string());
85    }
86    Ok(out)
87}
88
89/// Load a suite from a file path.
90pub fn load_suite(path: &Path) -> Result<Vec<LocomoSample>, String> {
91    let raw =
92        std::fs::read_to_string(path).map_err(|e| format!("reading {}: {e}", path.display()))?;
93    parse_suite(&raw)
94}
95
96/// The committed reference suite (real, verifiable facts; every gold answer is
97/// grounded in a turn and objectively true).
98pub const REFERENCE_SUITE: &str = include_str!("../../../data/locomo/reference-suite.ndjson");
99
100/// Parse the bundled reference suite. Panics only if the committed fixture is
101/// malformed, which a unit test guards against.
102pub fn reference_samples() -> Vec<LocomoSample> {
103    parse_suite(REFERENCE_SUITE).expect("bundled reference suite must be valid")
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn reference_suite_parses_and_is_grounded() {
112        let samples = reference_samples();
113        assert!(!samples.is_empty());
114        for s in &samples {
115            assert!(!s.sessions.is_empty(), "{} has no sessions", s.id);
116            assert!(!s.qa.is_empty(), "{} has no QA", s.id);
117            let transcript = s.transcript().to_lowercase();
118            // Every gold answer must be grounded somewhere in the transcript.
119            for qa in &s.qa {
120                assert!(!qa.answers.is_empty(), "QA without answers in {}", s.id);
121                let grounded = qa
122                    .answers
123                    .iter()
124                    .any(|a| transcript.contains(&a.to_lowercase()));
125                assert!(
126                    grounded,
127                    "answer for '{}' not grounded in transcript of {}",
128                    qa.question, s.id
129                );
130            }
131        }
132    }
133
134    #[test]
135    fn parses_json_array_form() {
136        let raw = r#"[{"id":"x","sessions":[{"session_id":"s","turns":[{"speaker":"A","text":"hi"}]}],"qa":[{"question":"q","answers":["hi"]}]}]"#;
137        let s = parse_suite(raw).unwrap();
138        assert_eq!(s.len(), 1);
139        assert_eq!(s[0].qa[0].category, 1, "default category applied");
140    }
141
142    #[test]
143    fn empty_suite_errors() {
144        assert!(parse_suite("# only a comment\n").is_err());
145    }
146}
lean_ctx/core/locomo/dataset.rs

lean_ctx/core/locomo/
dataset.rs
