1use std::collections::BTreeMap;
7
8use crate::extract::{admit_seat_write, SeatWrite};
9
10pub const COMPETENCIES: &[&str] = &[
12 "Accurate_Retrieval",
13 "Test_Time_Learning",
14 "Long_Range_Understanding",
15 "Conflict_Resolution",
16];
17
18#[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#[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 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 #[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}