lean_ctx/core/locomo/
runner.rs

1//! Run a LoCoMo sample through lean-ctx memory: ingest every turn as a knowledge
2//! fact, then for each question recall the top-k memories and score the recalled
3//! context against the gold answers (#291).
4//!
5//! This measures what lean-ctx actually does — *retrieval recall*: did the
6//! answer-bearing turn get surfaced, and at what token cost versus dumping the
7//! whole transcript. It is deliberately model-free, so results are deterministic.
8
9use std::path::Path;
10
11use crate::core::eval_ab::scorers::{qa_contains, qa_exact_match, qa_f1};
12use crate::core::knowledge::ProjectKnowledge;
13use crate::core::memory_policy::MemoryPolicy;
14use crate::core::tokens::count_tokens;
15
16use super::dataset::LocomoSample;
17
18/// Outcome of scoring a single question.
19#[derive(Debug, Clone)]
20pub struct QaResult {
21    pub category: u8,
22    pub f1: f64,
23    pub exact_match: bool,
24    /// A gold answer is a substring of the recalled context (the key recall signal).
25    pub contained: bool,
26    pub recall_tokens: usize,
27}
28
29/// Outcome of one sample.
30#[derive(Debug, Clone)]
31pub struct SampleResult {
32    pub id: String,
33    pub qa: Vec<QaResult>,
34    pub transcript_tokens: usize,
35}
36
37/// Ingest a sample's turns into an isolated knowledge store rooted at
38/// `project_root`, then recall + score each question with `top_k` memories.
39///
40/// `project_root` must be unique per sample so the knowledge hashes don't collide.
41/// The caller is responsible for pointing `LEAN_CTX_DATA_DIR` at a throwaway dir.
42pub fn run_sample(sample: &LocomoSample, project_root: &Path, top_k: usize) -> SampleResult {
43    let root = project_root.to_string_lossy().to_string();
44    let policy = MemoryPolicy::default();
45
46    // 1. Ingest every turn as a memory under a fresh store.
47    let _ = ProjectKnowledge::mutate_locked(&root, |k| {
48        for (si, session) in sample.sessions.iter().enumerate() {
49            let session_id = if session.session_id.is_empty() {
50                format!("s{si}")
51            } else {
52                session.session_id.clone()
53            };
54            for (ti, turn) in session.turns.iter().enumerate() {
55                let key = format!("{}-{}-{}", sample.id, session_id, ti);
56                let value = format!("{}: {}", turn.speaker, turn.text);
57                k.remember("conversation", &key, &value, &session_id, 0.9, &policy);
58            }
59        }
60    });
61
62    let transcript_tokens = count_tokens(&sample.transcript());
63
64    // 2. Recall + score each question against the production recall path.
65    let mut knowledge = ProjectKnowledge::load_or_create(&root);
66    let mut qa = Vec::with_capacity(sample.qa.len());
67    for item in &sample.qa {
68        let (hits, _) = knowledge.recall_for_output(&item.question, top_k);
69        let context = hits
70            .iter()
71            .map(|f| f.value.as_str())
72            .collect::<Vec<_>>()
73            .join("\n");
74        let recall_tokens = count_tokens(&context);
75
76        // Containment is measured over the whole recalled context (did any top-k
77        // memory carry the answer = retrieval recall). F1 / EM are measured against
78        // the *best single* recalled memory, so a short gold answer isn't penalised
79        // for the size of the surrounding context.
80        let mut best_f1 = 0.0f64;
81        let mut em = false;
82        let mut contained = false;
83        for gold in &item.answers {
84            contained |= qa_contains(&context, gold);
85            for hit in &hits {
86                best_f1 = best_f1.max(qa_f1(&hit.value, gold));
87                em |= qa_exact_match(&hit.value, gold);
88            }
89        }
90        qa.push(QaResult {
91            category: item.category,
92            f1: best_f1,
93            exact_match: em,
94            contained,
95            recall_tokens,
96        });
97    }
98
99    SampleResult {
100        id: sample.id.clone(),
101        qa,
102        transcript_tokens,
103    }
104}
105
106/// Run every sample under per-sample subdirectories of `workspace`.
107pub fn run_suite(samples: &[LocomoSample], workspace: &Path, top_k: usize) -> Vec<SampleResult> {
108    samples
109        .iter()
110        .enumerate()
111        .map(|(i, sample)| {
112            let proj = workspace.join(format!("sample-{i}"));
113            let _ = std::fs::create_dir_all(&proj);
114            run_sample(sample, &proj, top_k)
115        })
116        .collect()
117}
lean_ctx/core/locomo/runner.rs

lean_ctx/core/locomo/
runner.rs
