Skip to main content

packset_core/
mab.rs

1//! MemoryAgentBench write protocol: Remember / Prefer / Accept only.
2//!
3//! The four competencies are always named. A raw record is a refusal, not
4//! a zero. Hit rates are only over admitted writes.
5
6use std::collections::BTreeMap;
7
8use crate::extract::{admit_seat_write, SeatWrite};
9
10/// The benchmark's four competencies, in the paper's order.
11pub const COMPETENCIES: &[&str] = &[
12    "Accurate_Retrieval",
13    "Test_Time_Learning",
14    "Long_Range_Understanding",
15    "Conflict_Resolution",
16];
17
18/// One competency's asked/hit counts. `asked == 0` means the split was
19/// not run, not that the system scored zero.
20#[derive(Debug, Clone, Default, PartialEq, Eq)]
21pub struct Competency {
22    pub asked: usize,
23    pub hit: usize,
24}
25
26impl Competency {
27    #[must_use]
28    pub fn rate(&self) -> Option<f64> {
29        if self.asked == 0 {
30            None
31        } else {
32            Some(self.hit as f64 / self.asked as f64)
33        }
34    }
35}
36
37/// Honest report: refusals are counted, not scored as misses.
38#[derive(Debug, Clone, Default, PartialEq, Eq)]
39pub struct ProtocolReport {
40    pub admitted: usize,
41    pub refused: usize,
42    pub by_competency: BTreeMap<String, Competency>,
43}
44
45impl ProtocolReport {
46    #[must_use]
47    pub fn new() -> Self {
48        let mut by_competency = BTreeMap::new();
49        for name in COMPETENCIES {
50            by_competency.insert((*name).to_string(), Competency::default());
51        }
52        Self {
53            admitted: 0,
54            refused: 0,
55            by_competency,
56        }
57    }
58
59    /// Attempt to ingest one line. Raw context refuses.
60    pub fn ingest(&mut self, text: &str) -> Option<SeatWrite> {
61        match admit_seat_write(text) {
62            Some(w) => {
63                self.admitted += 1;
64                Some(w)
65            }
66            None => {
67                self.refused += 1;
68                None
69            }
70        }
71    }
72
73    pub fn mark(&mut self, competency: &str, hit: bool) {
74        let row = self
75            .by_competency
76            .entry(competency.to_string())
77            .or_default();
78        row.asked += 1;
79        if hit {
80            row.hit += 1;
81        }
82    }
83
84    #[must_use]
85    pub fn refusal_rate(&self) -> Option<f64> {
86        let n = self.admitted + self.refused;
87        if n == 0 {
88            None
89        } else {
90            Some(self.refused as f64 / n as f64)
91        }
92    }
93
94    /// Markdown table. Empty competencies print `—`, not 0.000.
95    #[must_use]
96    pub fn table(&self) -> String {
97        let mut s = String::from("| competency | asked | hit | rate |\n|---|---:|---:|---:|\n");
98        for name in COMPETENCIES {
99            let row = self.by_competency.get(*name).cloned().unwrap_or_default();
100            match row.rate() {
101                Some(r) => s.push_str(&format!(
102                    "| {name} | {} | {} | {:.3} |\n",
103                    row.asked, row.hit, r
104                )),
105                None => s.push_str(&format!("| {name} | 0 | — | — |\n")),
106            }
107        }
108        let n = self.admitted + self.refused;
109        let refuse = self
110            .refusal_rate()
111            .map(|r| format!("{r:.3}"))
112            .unwrap_or_else(|| "—".into());
113        s.push_str(&format!(
114            "| refusal | {n} | {} | {refuse} |\n",
115            self.refused
116        ));
117        s
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn raw_context_is_a_refusal_not_a_zero() {
127        let mut p = ProtocolReport::new();
128        assert!(p
129            .ingest("The user lives in Berlin and likes tea.")
130            .is_none());
131        assert_eq!(p.refused, 1);
132        assert_eq!(p.admitted, 0);
133        assert!(p.ingest("Remember: pin the review set").is_some());
134        assert_eq!(p.admitted, 1);
135        assert_eq!(p.by_competency.len(), 4);
136        assert!(p.by_competency["Test_Time_Learning"].rate().is_none());
137        let table = p.table();
138        assert!(
139            table.contains("| Test_Time_Learning | 0 | — | — |"),
140            "{table}"
141        );
142        assert!(table.contains("| refusal |"), "{table}");
143    }
144
145    #[test]
146    fn a_hit_is_only_counted_on_an_admitted_write() {
147        let mut p = ProtocolReport::new();
148        p.ingest("The user lives in Berlin.");
149        p.mark("Accurate_Retrieval", false);
150        p.ingest("Remember: Berlin is the capital of Germany now");
151        p.mark("Accurate_Retrieval", true);
152        let row = &p.by_competency["Accurate_Retrieval"];
153        assert_eq!(row.asked, 2);
154        assert_eq!(row.hit, 1);
155    }
156}