1use anyhow::Result;
15use serde::Serialize;
16
17use crate::server::Bridge;
18use crate::wire::{Row, ScoredAnswer};
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum Perturbation {
23 Reversed,
25}
26
27impl Perturbation {
28 pub const ALL: [Perturbation; 1] = [Perturbation::Reversed];
30
31 pub fn name(self) -> &'static str {
33 match self {
34 Perturbation::Reversed => "reversed-options",
35 }
36 }
37
38 pub fn apply(self, row: &Row) -> Row {
40 let mut changed = row.clone();
41 match self {
42 Perturbation::Reversed => changed.options.reverse(),
43 }
44 changed
45 }
46}
47
48#[derive(Debug, Serialize)]
50pub struct FlipRow {
51 pub id: String,
53 pub perturbation: String,
55 pub baseline: String,
57 pub perturbed: String,
59 pub flipped: bool,
61 pub baseline_probability: f64,
63 pub perturbed_probability: f64,
65}
66
67#[derive(Debug, Serialize)]
69pub struct FlipReport {
70 pub rows: usize,
72 pub flips: usize,
74 pub flip_rate: f64,
76 pub results: Vec<FlipRow>,
78}
79
80pub async fn run(bridge: &Bridge, rows: &[Row]) -> Result<FlipReport> {
82 let mut results = Vec::new();
83 for row in rows {
84 let baseline = bridge.score_row(row).await?;
85 let winner = argmax(&baseline.answer)?;
86 for perturbation in Perturbation::ALL {
87 let perturbed = bridge.score_row(&perturbation.apply(row)).await?;
88 let perturbed_winner = argmax(&perturbed.answer)?;
89 results.push(FlipRow {
90 id: row.id.clone(),
91 perturbation: perturbation.name().to_string(),
92 baseline: winner.clone(),
93 perturbed: perturbed_winner.clone(),
94 flipped: winner != perturbed_winner,
95 baseline_probability: probability_of(&baseline.answer, &winner)?,
96 perturbed_probability: probability_of(&perturbed.answer, &winner)?,
97 });
98 }
99 }
100 let flips = results.iter().filter(|row| row.flipped).count();
101 Ok(FlipReport {
102 rows: results.len(),
103 flips,
104 flip_rate: if results.is_empty() {
105 0.0
106 } else {
107 flips as f64 / results.len() as f64
108 },
109 results,
110 })
111}
112
113fn argmax(answer: &ScoredAnswer) -> Result<String> {
115 let index = answer
116 .probabilities
117 .iter()
118 .enumerate()
119 .max_by(|left, right| left.1.total_cmp(right.1))
120 .map(|(index, _)| index)
121 .ok_or_else(|| anyhow::anyhow!("a scored answer carries no probabilities"))?;
122 Ok(answer.option_ids[index].clone())
123}
124
125fn probability_of(answer: &ScoredAnswer, option: &str) -> Result<f64> {
127 let index = answer
128 .option_ids
129 .iter()
130 .position(|id| id == option)
131 .ok_or_else(|| anyhow::anyhow!("option {option:?} is missing from the scored answer"))?;
132 Ok(answer.probabilities[index])
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138 use serde_json::json;
139
140 fn row() -> Row {
141 serde_json::from_value(json!({
142 "id": "r1",
143 "state": "The customer was charged twice.",
144 "question": "Which department?",
145 "options": [
146 {"id": "billing", "description": "charges"},
147 {"id": "sales", "description": "purchases"},
148 ],
149 }))
150 .expect("a well-formed row")
151 }
152
153 #[test]
154 fn reversing_keeps_the_same_options_in_the_opposite_order() {
155 let original = row();
156 let reversed = Perturbation::Reversed.apply(&original);
157 assert_eq!(
158 reversed
159 .options
160 .iter()
161 .map(|option| option.id.as_str())
162 .collect::<Vec<_>>(),
163 vec!["sales", "billing"]
164 );
165 assert_eq!(original.options[0].id, "billing");
167 }
168
169 #[test]
170 fn the_winner_is_the_largest_probability_wherever_it_sits() {
171 let answer = ScoredAnswer {
172 id: "r1".to_string(),
173 option_ids: vec!["sales".to_string(), "billing".to_string()],
174 probabilities: vec![0.2, 0.8],
175 input_tokens: 0,
176 prompt_version: None,
177 };
178 assert_eq!(argmax(&answer).unwrap(), "billing");
179 assert_eq!(probability_of(&answer, "billing").unwrap(), 0.8);
180 assert!(probability_of(&answer, "missing").is_err());
181 }
182}