Skip to main content

jev_repl/
evaluate.rs

1//! Scoring a rubric against states someone has already labelled.
2//!
3//! One call tells you what the model said about one state; that is what `jev run` is for. It does
4//! not tell you where to put the threshold, or how confident a choice has to be before a script
5//! may act on it — the README leaves both calls to you, and this is the module that turns them
6//! into a table. Feed it a page and a file of labelled states and it reports what the rubric got
7//! right, at every threshold worth trying.
8//!
9//! Everything here is pure: cases come in as text, the answers come from a function you pass, and
10//! the report goes out as lines and JSON. Reading files, hashing request bodies and talking to the
11//! API belong to the caller.
12
13use std::future::Future;
14use std::sync::Arc;
15
16use ratatui::style::Style;
17use ratatui::text::{Line, Span};
18use serde_json::{Value, json};
19use tokio::sync::Semaphore;
20use tokio::task::JoinSet;
21use typesafe::{Answer, Choice, Question, Usage};
22
23use crate::cost::{self, Cost, Rates};
24use crate::format::{BAD, CHOICE, bold, color_for, dim, text_of};
25use crate::headless::Answered;
26use crate::session::{self, Session};
27
28/// One labelled state: what to judge, and what the rubric should say about it.
29#[derive(Debug, Clone, PartialEq)]
30pub struct Case {
31    /// 1-based line in the cases file, for messages.
32    pub line: usize,
33    pub id: Option<String>,
34    pub state: Value,
35    /// Question name → expectation, in the order the file gave them, already checked against the
36    /// session's questions. A `Vec` and not a map, because that is the shape `Session` uses for
37    /// questions and it saves a dependency.
38    pub expect: Vec<(String, Expectation)>,
39}
40
41/// What one question is expected to answer, in the shape its kind is scored in.
42#[derive(Debug, Clone, PartialEq)]
43pub enum Expectation {
44    Noul { yes: bool },
45    Choice { label: String },
46    Score { level: usize },
47}
48
49impl Expectation {
50    /// The wire `type` of the answer this expectation can be compared with.
51    pub fn kind(&self) -> &'static str {
52        match self {
53            Expectation::Noul { .. } => "noul",
54            Expectation::Choice { .. } => "choice",
55            Expectation::Score { .. } => "score",
56        }
57    }
58}
59
60/// Parse JSON Lines into cases, checking every expectation against `session`.
61///
62/// Blank lines are skipped and everything else has to be a case, because a file of labels is worth
63/// nothing if a typo silently drops a row. The line number travels with the case: it is what the
64/// report names a case by, so a bad row is found by the same number that reported it.
65pub fn parse_cases(text: &str, session: &Session) -> Result<Vec<Case>, String> {
66    let mut cases = Vec::new();
67    for (i, raw) in text.split('\n').enumerate() {
68        let raw = raw.trim();
69        if raw.is_empty() {
70            continue;
71        }
72        let line = i + 1;
73        let one = parse_case(raw, line, session).map_err(|e| format!("cases line {line}: {e}"))?;
74        cases.push(one);
75    }
76    if cases.is_empty() {
77        return Err("the cases file holds no cases.".to_owned());
78    }
79    Ok(cases)
80}
81
82fn parse_case(text: &str, line: usize, session: &Session) -> Result<Case, String> {
83    let value: Value = serde_json::from_str(text).map_err(|e| format!("not valid JSON: {e}"))?;
84    let object = value
85        .as_object()
86        .ok_or("expected a JSON object with `state` and `expect`.")?;
87
88    let id = match object.get("id") {
89        None => None,
90        Some(Value::String(s)) => Some(s.clone()),
91        Some(_) => return Err("`id` must be a string.".to_owned()),
92    };
93
94    let state = object
95        .get("state")
96        .ok_or("missing `state`: a case has to say what to judge.")?;
97    if session::is_empty_value(state) {
98        return Err("the `state` is empty: there is nothing to judge.".to_owned());
99    }
100
101    let wanted = object
102        .get("expect")
103        .ok_or("missing `expect`: a case has to say what the answer is.")?;
104    let wanted = wanted
105        .as_object()
106        .filter(|map| !map.is_empty())
107        .ok_or("`expect` has to name at least one question.")?;
108
109    let mut expect = Vec::with_capacity(wanted.len());
110    for (name, value) in wanted {
111        let question = session
112            .questions
113            .iter()
114            .find(|(n, _)| n == name)
115            .map(|(_, q)| q)
116            .ok_or_else(|| format!("no question named {name:?} on the page."))?;
117        expect.push((name.clone(), expected(name, question, value)?));
118    }
119    Ok(Case {
120        line,
121        id,
122        state: state.clone(),
123        expect,
124    })
125}
126
127/// Check one expected value against the question it names, and store it the way it is scored.
128fn expected(name: &str, question: &Question, value: &Value) -> Result<Expectation, String> {
129    match question {
130        Question::Noul(_) => match value {
131            Value::Bool(yes) => Ok(Expectation::Noul { yes: *yes }),
132            other => Err(format!(
133                "{name} is a noul: expected true or false, got {other}."
134            )),
135        },
136        Question::Choice(q) => {
137            let labels: Vec<&str> = q.criteria.keys().map(String::as_str).collect();
138            match value.as_str() {
139                Some(label) if labels.contains(&label) => Ok(Expectation::Choice {
140                    label: label.to_owned(),
141                }),
142                _ => Err(format!(
143                    "{name} is a choice between {}; got {value}.",
144                    labels.join(", ")
145                )),
146            }
147        }
148        Question::Score(q) => {
149            let top = q.criteria.len().saturating_sub(1);
150            if let Value::Number(number) = value {
151                let n = number.as_f64().unwrap_or(f64::NAN);
152                if n.fract() == 0.0 && (0.0..=top as f64).contains(&n) {
153                    return Ok(Expectation::Score { level: n as usize });
154                }
155                return Err(format!(
156                    "{name} is a score: expected a level from 0 to {top}, got {number}."
157                ));
158            }
159            // A level's own text reads better in a cases file than its index does; the first wins.
160            let wanted = text_of(value);
161            match q.criteria.iter().position(|level| text_of(level) == wanted) {
162                Some(at) => Ok(Expectation::Score { level: at }),
163                None => Err(format!(
164                    "{name} is a score: expected a level from 0 to {top}, or one of its levels; got {value}."
165                )),
166            }
167        }
168        // A hand-built question object has no shape to score against, and neither has a kind this
169        // build does not know.
170        _ => Err(format!(
171            "{name} is a raw question: raw questions cannot be scored."
172        )),
173    }
174}
175
176/// The session as one case sends it: the page's questions, the case's state.
177pub fn with_state(session: &Session, state: Value) -> Session {
178    Session {
179        state,
180        questions: session.questions.clone(),
181        model: session.model.clone(),
182    }
183}
184
185/// What one case's request came back as.
186#[derive(Debug, Clone)]
187pub enum Outcome {
188    Ok {
189        answers: Vec<Answered>,
190        usage: Option<Usage>,
191    },
192    Failed {
193        error: String,
194    },
195}
196
197/// Send every case through `ask`, at most `concurrency` at a time; results are in case order.
198///
199/// Every case is spawned at once and a semaphore decides how many are in the air, so a slow case
200/// holds up nothing but itself, and each result is written to its own slot: a file of a thousand
201/// labels keeps its order however the calls come back.
202pub async fn run<F, Fut>(
203    session: &Session,
204    cases: &[Case],
205    ask: F,
206    concurrency: usize,
207) -> Vec<Outcome>
208where
209    F: Fn(Session) -> Fut + Send + Sync + Clone + 'static,
210    Fut: Future<Output = Outcome> + Send + 'static,
211{
212    let permits = Arc::new(Semaphore::new(concurrency.max(1)));
213    let mut workers = JoinSet::new();
214    for (at, one) in cases.iter().enumerate() {
215        let session = with_state(session, one.state.clone());
216        let ask = ask.clone();
217        let permits = Arc::clone(&permits);
218        workers.spawn(async move {
219            let _permit = permits.acquire_owned().await;
220            (at, ask(session).await)
221        });
222    }
223
224    let mut outcomes: Vec<Option<Outcome>> = vec![None; cases.len()];
225    while let Some(joined) = workers.join_next().await {
226        // A worker that panicked leaves its slot empty; the run is not lost to one case.
227        if let Ok((at, outcome)) = joined {
228            outcomes[at] = Some(outcome);
229        }
230    }
231    outcomes
232        .into_iter()
233        .map(|outcome| {
234            outcome.unwrap_or_else(|| Outcome::Failed {
235                error: "nothing was sent for this case.".to_owned(),
236            })
237        })
238        .collect()
239}
240
241/// One row of a noul's threshold sweep: the confusion counts, and what they come to.
242#[derive(Debug, Clone, PartialEq)]
243pub struct SweepRow {
244    pub threshold: f64,
245    pub tp: usize,
246    pub fp: usize,
247    pub r#fn: usize,
248    pub tn: usize,
249    pub accuracy: f64,
250    /// `None` when nothing was predicted a yes, because a rate over nothing is not zero.
251    pub precision: Option<f64>,
252    /// `None` when nothing was expected to be a yes.
253    pub recall: Option<f64>,
254    pub f1: f64,
255}
256
257/// One cut of the confidence gate: how much of the set survives it, and how right it is.
258#[derive(Debug, Clone, PartialEq)]
259pub struct GateRow {
260    pub confidence: f64,
261    pub coverage: f64,
262    /// `None` when no case is confident enough to be counted.
263    pub accuracy: Option<f64>,
264}
265
266/// The threshold that scored best, and what it scored.
267#[derive(Debug, Clone, Copy, PartialEq)]
268pub struct Best {
269    pub threshold: f64,
270    pub f1: f64,
271}
272
273/// What one question scored, in the numbers its kind is judged by.
274#[derive(Debug, Clone, PartialEq)]
275pub enum QuestionReport {
276    Noul {
277        name: String,
278        cases: usize,
279        /// Mean squared error of the probability itself, threshold or no threshold.
280        brier: f64,
281        /// Accuracy at the threshold this run was asked to use.
282        accuracy: f64,
283        best: Best,
284        sweep: Vec<SweepRow>,
285    },
286    Choice {
287        name: String,
288        cases: usize,
289        accuracy: f64,
290        /// The page's options, plus `other` when the model answered something else.
291        labels: Vec<String>,
292        /// Rows expected, columns predicted.
293        confusion: Vec<Vec<usize>>,
294        gate: Vec<GateRow>,
295    },
296    Score {
297        name: String,
298        cases: usize,
299        exact: f64,
300        within_one: f64,
301        mae: f64,
302        gate: Vec<GateRow>,
303    },
304}
305
306impl QuestionReport {
307    pub fn name(&self) -> &str {
308        match self {
309            QuestionReport::Noul { name, .. }
310            | QuestionReport::Choice { name, .. }
311            | QuestionReport::Score { name, .. } => name,
312        }
313    }
314
315    /// The wire `type` of the question this reports on.
316    pub fn kind(&self) -> &'static str {
317        match self {
318            QuestionReport::Noul { .. } => "noul",
319            QuestionReport::Choice { .. } => "choice",
320            QuestionReport::Score { .. } => "score",
321        }
322    }
323
324    pub fn cases(&self) -> usize {
325        match self {
326            QuestionReport::Noul { cases, .. }
327            | QuestionReport::Choice { cases, .. }
328            | QuestionReport::Score { cases, .. } => *cases,
329        }
330    }
331
332    /// The accuracy `--min-accuracy` holds a question to: exact agreement at the chosen threshold.
333    pub fn accuracy_of(&self) -> f64 {
334        match self {
335            QuestionReport::Noul { accuracy, .. } | QuestionReport::Choice { accuracy, .. } => {
336                *accuracy
337            }
338            QuestionReport::Score { exact, .. } => *exact,
339        }
340    }
341}
342
343/// A case that never produced a full set of answers, and why.
344#[derive(Debug, Clone, PartialEq)]
345pub struct CaseError {
346    pub case: usize,
347    pub id: Option<String>,
348    pub message: String,
349}
350
351/// The tokens the run spent, counted when the API counted them and estimated when it did not.
352#[derive(Debug, Clone, PartialEq)]
353pub struct ReportUsage {
354    pub input_tokens: u64,
355    pub output_tokens: u64,
356    pub estimated: bool,
357    pub cost: Option<Cost>,
358}
359
360/// Everything the run found out, with the numbers unrounded.
361#[derive(Debug, Clone, PartialEq)]
362pub struct Report {
363    pub model: String,
364    pub threshold: f64,
365    pub cases: usize,
366    pub answered: usize,
367    pub errors: Vec<CaseError>,
368    pub questions: Vec<QuestionReport>,
369    pub usage: ReportUsage,
370}
371
372/// What the run was asked for, which the report repeats back.
373#[derive(Debug, Clone, Copy)]
374pub struct ReportOptions<'a> {
375    pub model: &'a str,
376    pub threshold: f64,
377    pub rates: Option<Rates>,
378}
379
380/// The thresholds a sweep always covers; the chosen one joins them when it is not one of these.
381const SWEEP: [f64; 9] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9];
382
383/// The cuts the confidence gate is read at.
384const CUTS: [f64; 5] = [0.0, 0.2, 0.4, 0.6, 0.8];
385
386/// One case that answered everything it was labelled for.
387struct Scored<'a> {
388    expect: &'a [(String, Expectation)],
389    answers: Vec<(String, Answer)>,
390    usage: Option<Usage>,
391}
392
393impl Scored<'_> {
394    fn answer(&self, name: &str) -> Option<&Answer> {
395        self.answers.iter().find(|(n, _)| n == name).map(|(_, a)| a)
396    }
397
398    fn expects(&self, name: &str) -> Option<&Expectation> {
399        self.expect.iter().find(|(n, _)| n == name).map(|(_, e)| e)
400    }
401}
402
403/// Score the outcomes against the cases.
404///
405/// A case either answered everything it was labelled for or it counts as an error: a half-answered
406/// case would quietly skew whichever question it did answer, and a rubric is being judged here.
407pub fn report(
408    session: &Session,
409    cases: &[Case],
410    outcomes: &[Outcome],
411    options: ReportOptions<'_>,
412) -> Report {
413    let mut errors: Vec<CaseError> = Vec::new();
414    let mut scored: Vec<Scored<'_>> = Vec::new();
415    for (at, one) in cases.iter().enumerate() {
416        let mut failed = |message: String| {
417            errors.push(CaseError {
418                case: one.line,
419                id: one.id.clone(),
420                message,
421            });
422        };
423        let (answers, usage) = match outcomes.get(at) {
424            None => {
425                failed("nothing was sent for this case.".to_owned());
426                continue;
427            }
428            Some(Outcome::Failed { error }) => {
429                failed(error.clone());
430                continue;
431            }
432            Some(Outcome::Ok { answers, usage }) => (answers, usage),
433        };
434        let answers: Vec<(String, Answer)> = answers
435            .iter()
436            .filter_map(|(name, answer)| answer.clone().map(|a| (name.clone(), a)))
437            .collect();
438        if let Some(message) = unscorable(&one.expect, &answers) {
439            failed(message);
440            continue;
441        }
442        scored.push(Scored {
443            expect: &one.expect,
444            answers,
445            usage: usage.clone(),
446        });
447    }
448
449    let mut questions = Vec::new();
450    for (name, question) in &session.questions {
451        let rows: Vec<&Scored<'_>> = scored
452            .iter()
453            .filter(|one| one.expects(name).is_some())
454            .collect();
455        if rows.is_empty() {
456            continue;
457        }
458        match question {
459            Question::Noul(_) => questions.push(noul_report(name, &rows, options.threshold)),
460            Question::Choice(q) => questions.push(choice_report(name, q, &rows)),
461            Question::Score(_) => questions.push(score_report(name, &rows)),
462            _ => {}
463        }
464    }
465
466    let usage = usage_of(session, cases, &scored, options.model, options.rates);
467    Report {
468        model: options.model.to_owned(),
469        threshold: options.threshold,
470        cases: cases.len(),
471        answered: scored.len(),
472        errors,
473        questions,
474        usage,
475    }
476}
477
478/// Why this case cannot be scored, if it cannot: the first label that got no answer, or one whose
479/// answer came back as another kind.
480fn unscorable(expect: &[(String, Expectation)], answers: &[(String, Answer)]) -> Option<String> {
481    for (name, expectation) in expect {
482        match answers.iter().find(|(n, _)| n == name).map(|(_, a)| a) {
483            None => return Some(format!("no answer came back for {name}")),
484            Some(answer) if answer.kind() != expectation.kind() => {
485                return Some(format!(
486                    "{name} came back as a {}, not a {}",
487                    answer.kind(),
488                    expectation.kind()
489                ));
490            }
491            Some(_) => {}
492        }
493    }
494    None
495}
496
497/// Questions whose accuracy is below `bar`, for --min-accuracy.
498pub fn below_bar(report: &Report, bar: f64) -> Vec<(String, f64)> {
499    report
500        .questions
501        .iter()
502        .filter(|q| q.accuracy_of() < bar)
503        .map(|q| (q.name().to_owned(), q.accuracy_of()))
504        .collect()
505}
506
507fn noul_report(name: &str, rows: &[&Scored<'_>], threshold: f64) -> QuestionReport {
508    let points: Vec<(f64, bool)> = rows
509        .iter()
510        .map(|row| {
511            let p = match row.answer(name) {
512                Some(Answer::Noul(a)) => a.noul,
513                _ => 0.0,
514            };
515            let yes = matches!(row.expects(name), Some(Expectation::Noul { yes: true }));
516            (p, yes)
517        })
518        .collect();
519
520    let mut thresholds: Vec<f64> = SWEEP.to_vec();
521    if !thresholds.contains(&threshold) {
522        thresholds.push(threshold);
523        thresholds.sort_by(f64::total_cmp);
524    }
525    let sweep: Vec<SweepRow> = thresholds
526        .iter()
527        .map(|at| sweep_row(&points, *at))
528        .collect();
529    let accuracy = sweep
530        .iter()
531        .find(|row| row.threshold == threshold)
532        .map(|row| row.accuracy)
533        .unwrap_or(0.0);
534    // The sweep is in ascending order and the comparison is strict, so a tie keeps the lowest.
535    let mut best = Best {
536        threshold,
537        f1: f64::NEG_INFINITY,
538    };
539    for row in &sweep {
540        if row.f1 > best.f1 {
541            best = Best {
542                threshold: row.threshold,
543                f1: row.f1,
544            };
545        }
546    }
547    let brier = mean(
548        points
549            .iter()
550            .map(|(p, yes)| (p - if *yes { 1.0 } else { 0.0 }).powi(2)),
551    );
552    QuestionReport::Noul {
553        name: name.to_owned(),
554        cases: points.len(),
555        brier,
556        accuracy,
557        best,
558        sweep,
559    }
560}
561
562fn sweep_row(points: &[(f64, bool)], threshold: f64) -> SweepRow {
563    let (mut tp, mut fp, mut fneg, mut tn) = (0usize, 0usize, 0usize, 0usize);
564    for (p, yes) in points {
565        match (*p >= threshold, *yes) {
566            (true, true) => tp += 1,
567            (true, false) => fp += 1,
568            (false, true) => fneg += 1,
569            (false, false) => tn += 1,
570        }
571    }
572    let denominator = 2 * tp + fp + fneg;
573    SweepRow {
574        threshold,
575        tp,
576        fp,
577        r#fn: fneg,
578        tn,
579        accuracy: (tp + tn) as f64 / points.len() as f64,
580        precision: (tp + fp > 0).then(|| tp as f64 / (tp + fp) as f64),
581        recall: (tp + fneg > 0).then(|| tp as f64 / (tp + fneg) as f64),
582        f1: if denominator == 0 {
583            0.0
584        } else {
585            2.0 * tp as f64 / denominator as f64
586        },
587    }
588}
589
590fn choice_report(name: &str, question: &Choice, rows: &[&Scored<'_>]) -> QuestionReport {
591    let options: Vec<String> = question.criteria.keys().cloned().collect();
592    struct Point {
593        predicted: String,
594        expected: String,
595        confidence: f64,
596        right: bool,
597    }
598    let points: Vec<Point> = rows
599        .iter()
600        .map(|row| {
601            let (predicted, confidence) = match row.answer(name) {
602                Some(Answer::Choice(a)) => (a.choice.clone(), a.confidence),
603                _ => (String::new(), 0.0),
604            };
605            let expected = match row.expects(name) {
606                Some(Expectation::Choice { label }) => label.clone(),
607                _ => String::new(),
608            };
609            Point {
610                right: predicted == expected,
611                predicted,
612                expected,
613                confidence,
614            }
615        })
616        .collect();
617
618    // A label the page never offered still has to land somewhere, or the matrix loses cases.
619    let other = points
620        .iter()
621        .any(|point| !options.contains(&point.predicted));
622    let mut labels = options.clone();
623    if other {
624        labels.push("other".to_owned());
625    }
626    let confusion: Vec<Vec<usize>> = options
627        .iter()
628        .map(|expected| {
629            labels
630                .iter()
631                .enumerate()
632                .map(|(column, predicted)| {
633                    points
634                        .iter()
635                        .filter(|point| {
636                            &point.expected == expected
637                                && if other && column == labels.len() - 1 {
638                                    !options.contains(&point.predicted)
639                                } else {
640                                    &point.predicted == predicted
641                                }
642                        })
643                        .count()
644                })
645                .collect()
646        })
647        .collect();
648
649    QuestionReport::Choice {
650        name: name.to_owned(),
651        cases: points.len(),
652        accuracy: mean(points.iter().map(|point| f64::from(point.right))),
653        labels,
654        confusion,
655        gate: gate(points.iter().map(|point| (point.confidence, point.right))),
656    }
657}
658
659fn score_report(name: &str, rows: &[&Scored<'_>]) -> QuestionReport {
660    let points: Vec<(f64, i64)> = rows
661        .iter()
662        .map(|row| {
663            let (level, confidence) = match row.answer(name) {
664                Some(Answer::Score(a)) => (i64::from(a.rounded_level()), a.confidence),
665                _ => (0, 0.0),
666            };
667            let expected = match row.expects(name) {
668                Some(Expectation::Score { level }) => *level as i64,
669                _ => 0,
670            };
671            (confidence, (level - expected).abs())
672        })
673        .collect();
674    QuestionReport::Score {
675        name: name.to_owned(),
676        cases: points.len(),
677        exact: mean(points.iter().map(|(_, off)| f64::from(*off == 0))),
678        within_one: mean(points.iter().map(|(_, off)| f64::from(*off <= 1))),
679        mae: mean(points.iter().map(|(_, off)| *off as f64)),
680        gate: gate(points.iter().map(|(c, off)| (*c, *off == 0))),
681    }
682}
683
684/// Coverage and accuracy at each cut: what you buy by only acting on confident answers.
685fn gate(points: impl Iterator<Item = (f64, bool)>) -> Vec<GateRow> {
686    let points: Vec<(f64, bool)> = points.collect();
687    CUTS.iter()
688        .map(|confidence| {
689            let kept: Vec<bool> = points
690                .iter()
691                .filter(|(c, _)| c >= confidence)
692                .map(|(_, right)| *right)
693                .collect();
694            GateRow {
695                confidence: *confidence,
696                coverage: if points.is_empty() {
697                    0.0
698                } else {
699                    kept.len() as f64 / points.len() as f64
700                },
701                accuracy: (!kept.is_empty())
702                    .then(|| mean(kept.iter().map(|right| f64::from(*right)))),
703            }
704        })
705        .collect()
706}
707
708/// Counted tokens when every answered case carried them; the estimate, marked as one, otherwise.
709fn usage_of(
710    session: &Session,
711    cases: &[Case],
712    scored: &[Scored<'_>],
713    model: &str,
714    rates: Option<Rates>,
715) -> ReportUsage {
716    let mut input_tokens = 0u64;
717    let mut output_tokens = 0u64;
718    let mut counted = !scored.is_empty();
719    for one in scored {
720        match one
721            .usage
722            .as_ref()
723            .map(|u| (u.input_tokens, u.output_tokens))
724        {
725            Some((Some(input), Some(output))) => {
726                input_tokens += input;
727                output_tokens += output;
728            }
729            _ => {
730                counted = false;
731                break;
732            }
733        }
734    }
735    if !counted {
736        let estimate = preflight(session, cases, model, None);
737        input_tokens = estimate.input_tokens as u64;
738        output_tokens = estimate.output_tokens as u64;
739    }
740    ReportUsage {
741        input_tokens,
742        output_tokens,
743        estimated: !counted,
744        cost: rates.map(|rates| cost::price(input_tokens, output_tokens, rates)),
745    }
746}
747
748/// What a whole run would send, before any of it is sent.
749#[derive(Debug, Clone, PartialEq)]
750pub struct Preflight {
751    pub cases: usize,
752    pub input_tokens: usize,
753    pub output_tokens: usize,
754    pub cost: Option<Cost>,
755}
756
757/// The preflight estimate: tokens summed over every case, priced when rates are known.
758pub fn preflight(
759    session: &Session,
760    cases: &[Case],
761    model: &str,
762    rates: Option<Rates>,
763) -> Preflight {
764    let mut input_tokens = 0usize;
765    let mut output_tokens = 0usize;
766    for one in cases {
767        let estimate = cost::estimate(&with_state(session, one.state.clone()), model);
768        input_tokens += estimate.input_tokens;
769        output_tokens += estimate.output_tokens;
770    }
771    Preflight {
772        cases: cases.len(),
773        input_tokens,
774        output_tokens,
775        cost: rates.map(|rates| cost::price(input_tokens as u64, output_tokens as u64, rates)),
776    }
777}
778
779fn mean(values: impl Iterator<Item = f64>) -> f64 {
780    let mut sum = 0.0;
781    let mut count = 0usize;
782    for value in values {
783        sum += value;
784        count += 1;
785    }
786    if count == 0 { 0.0 } else { sum / count as f64 }
787}
788
789/// Two decimals, rounding a tie away from zero the way the TypeScript port's `toFixed(2)` does.
790///
791/// Rust rounds a tie to even, so `0.125` would print `0.12` here and `0.13` there; every rate and
792/// probability in a report goes through this so the two ports' reports can be compared byte for
793/// byte.
794pub fn two(x: f64) -> String {
795    format!("{:.2}", (x * 100.0).round() / 100.0)
796}
797
798/// The text report, as lines the terminal draws.
799///
800/// One block per question, in the page's order: what it scored, the sweep or the gate that says
801/// where to set the dial, and — for a choice — the matrix that says what it confuses with what.
802pub fn report_lines(report: &Report) -> Vec<Line<'static>> {
803    let mut out: Vec<Line<'static>> = Vec::new();
804    let width = report
805        .questions
806        .iter()
807        .map(|q| q.name().chars().count())
808        .max()
809        .unwrap_or(0);
810    for question in &report.questions {
811        if !out.is_empty() {
812            out.push(Line::default());
813        }
814        out.push(header_line(question, width));
815        match question {
816            QuestionReport::Noul { sweep, best, .. } => {
817                out.extend(sweep_lines(sweep, *best, report.threshold));
818            }
819            QuestionReport::Choice {
820                gate,
821                labels,
822                confusion,
823                ..
824            } => {
825                out.extend(gate_lines(gate, "accuracy"));
826                out.extend(confusion_lines(labels, confusion));
827            }
828            QuestionReport::Score { gate, .. } => out.extend(gate_lines(gate, "exact")),
829        }
830    }
831
832    if !report.errors.is_empty() {
833        if !out.is_empty() {
834            out.push(Line::default());
835        }
836        for failed in &report.errors {
837            out.extend(error_case_lines(failed));
838        }
839    }
840
841    if !out.is_empty() {
842        out.push(Line::default());
843    }
844    let errors = report.errors.len();
845    out.push(Line::from(vec![
846        Span::raw("  "),
847        bold(format!("{} case{}", report.cases, plural(report.cases))),
848        dim(format!(
849            " · {} answered · {errors} error{}",
850            report.answered,
851            plural(errors)
852        )),
853    ]));
854    out.push(usage_line(&report.usage));
855    out
856}
857
858fn header_line(question: &QuestionReport, width: usize) -> Line<'static> {
859    let count = format!("{} case{}", question.cases(), plural(question.cases()));
860    let summary = match question {
861        QuestionReport::Noul { brier, .. } => format!("{count} · Brier {}", two(*brier)),
862        QuestionReport::Choice { accuracy, .. } => format!("{count} · accuracy {}", two(*accuracy)),
863        QuestionReport::Score {
864            exact,
865            within_one,
866            mae,
867            ..
868        } => format!(
869            "{count} · exact {} · within one {} · mae {}",
870            two(*exact),
871            two(*within_one),
872            two(*mae)
873        ),
874    };
875    Line::from(vec![
876        Span::raw("  "),
877        bold(pad_end(question.name(), width)),
878        Span::raw("  "),
879        Span::styled(
880            pad_end(question.kind(), 8),
881            Style::new().fg(color_for(question.kind())),
882        ),
883        dim(summary),
884    ])
885}
886
887/// The sweep: what the threshold buys, row by row, with a `*` on the one this run used.
888fn sweep_lines(sweep: &[SweepRow], best: Best, threshold: f64) -> Vec<Line<'static>> {
889    let mut out = vec![Line::from(vec![
890        Span::raw("    "),
891        dim(pad_end("threshold", 12)),
892        dim(pad_end("acc", 6)),
893        dim(pad_end("prec", 7)),
894        dim(pad_end("rec", 7)),
895        dim("f1"),
896    ])];
897    for row in sweep {
898        let chosen = row.threshold == threshold;
899        let at = pad_end(
900            &format!("{}{}", two(row.threshold), if chosen { " *" } else { "" }),
901            12,
902        );
903        out.push(Line::from(vec![
904            Span::raw("    "),
905            if chosen { bold(at) } else { Span::raw(at) },
906            Span::raw(pad_end(&two(row.accuracy), 6)),
907            Span::raw(pad_end(&rate(row.precision), 7)),
908            Span::raw(pad_end(&rate(row.recall), 7)),
909            Span::raw(two(row.f1)),
910        ]));
911    }
912    out.push(Line::from(vec![
913        Span::raw("    "),
914        dim(format!("best f1 at {}", two(best.threshold))),
915    ]));
916    out
917}
918
919fn gate_lines(gate: &[GateRow], accuracy: &str) -> Vec<Line<'static>> {
920    let mut out = vec![Line::from(vec![
921        Span::raw("    "),
922        dim(pad_end("confidence ≥", 15)),
923        dim(pad_end("coverage", 10)),
924        dim(accuracy.to_owned()),
925    ])];
926    for row in gate {
927        out.push(Line::from(vec![
928            Span::raw("    "),
929            Span::raw(pad_end(&two(row.confidence), 15)),
930            Span::raw(pad_end(&two(row.coverage), 10)),
931            Span::raw(rate(row.accuracy)),
932        ]));
933    }
934    out
935}
936
937/// The matrix, which is where a rubric's real confusions show: what it calls what.
938fn confusion_lines(labels: &[String], confusion: &[Vec<usize>]) -> Vec<Line<'static>> {
939    let counts: Vec<usize> = confusion
940        .iter()
941        .flatten()
942        .map(|n| n.to_string().len())
943        .collect();
944    let column = |label: &str| -> usize {
945        counts
946            .iter()
947            .copied()
948            .chain([label.chars().count(), 1])
949            .max()
950            .unwrap_or(1)
951            + 2
952    };
953    let row_width = confusion
954        .iter()
955        .enumerate()
956        .map(|(at, _)| labels[at].chars().count())
957        .max()
958        .unwrap_or(0)
959        + 3;
960
961    let heading: String = labels
962        .iter()
963        .map(|label| pad_end(label, column(label)))
964        .collect();
965    let mut out = vec![
966        Line::from(vec![
967            Span::raw("    "),
968            dim("confusion, rows expected, columns predicted"),
969        ]),
970        Line::from(vec![
971            Span::raw(format!("    {}", " ".repeat(row_width))),
972            dim(heading.trim_end().to_owned()),
973        ]),
974    ];
975    for (at, row) in confusion.iter().enumerate() {
976        let cells: String = row
977            .iter()
978            .enumerate()
979            .map(|(column2, count)| pad_end(&count.to_string(), column(&labels[column2])))
980            .collect();
981        out.push(Line::from(vec![
982            Span::raw("    "),
983            Span::styled(pad_end(&labels[at], row_width), Style::new().fg(CHOICE)),
984            Span::raw(cells.trim_end().to_owned()),
985        ]));
986    }
987    out
988}
989
990fn error_case_lines(failed: &CaseError) -> Vec<Line<'static>> {
991    let name = match &failed.id {
992        Some(id) => format!("case {} ({id})", failed.case),
993        None => format!("case {}", failed.case),
994    };
995    let mut parts = failed.message.split('\n');
996    let first = parts.next().unwrap_or("").trim().to_owned();
997    let mut out = vec![Line::from(vec![
998        Span::raw("  "),
999        Span::styled(format!("{name}: "), Style::new().fg(BAD)),
1000        Span::raw(first),
1001    ])];
1002    for more in parts {
1003        out.push(Line::from(vec![
1004            Span::raw("    "),
1005            dim(more.trim().to_owned()),
1006        ]));
1007    }
1008    out
1009}
1010
1011fn usage_line(usage: &ReportUsage) -> Line<'static> {
1012    let money = match usage.cost {
1013        Some(cost) => format!(" · {}", cost::usd(cost.total)),
1014        None => String::new(),
1015    };
1016    let tokens = format!(
1017        "{} in / {} out tokens{money}",
1018        usage.input_tokens, usage.output_tokens
1019    );
1020    if usage.estimated {
1021        Line::from(vec![
1022            Span::raw("  "),
1023            dim(format!("≈ {tokens} — estimated, nothing was counted")),
1024        ])
1025    } else {
1026        Line::from(vec![Span::raw("  "), dim(tokens)])
1027    }
1028}
1029
1030/// The JSON report, ready for `to_string_pretty`. Numbers keep their precision; what is undefined
1031/// is null.
1032pub fn report_json(report: &Report) -> Value {
1033    let mut questions = serde_json::Map::new();
1034    for question in &report.questions {
1035        questions.insert(question.name().to_owned(), question_json(question));
1036    }
1037    let mut usage = serde_json::Map::new();
1038    usage.insert("inputTokens".to_owned(), json!(report.usage.input_tokens));
1039    usage.insert("outputTokens".to_owned(), json!(report.usage.output_tokens));
1040    usage.insert("estimated".to_owned(), json!(report.usage.estimated));
1041    if let Some(cost) = report.usage.cost {
1042        usage.insert("cost".to_owned(), number(cost.total));
1043    }
1044    let errors: Vec<Value> = report
1045        .errors
1046        .iter()
1047        .map(|failed| {
1048            let mut out = serde_json::Map::new();
1049            out.insert("case".to_owned(), json!(failed.case));
1050            if let Some(id) = &failed.id {
1051                out.insert("id".to_owned(), json!(id));
1052            }
1053            out.insert("message".to_owned(), json!(failed.message));
1054            Value::Object(out)
1055        })
1056        .collect();
1057    json!({
1058        "model": report.model,
1059        "threshold": number(report.threshold),
1060        "cases": report.cases,
1061        "answered": report.answered,
1062        "errors": errors,
1063        "questions": Value::Object(questions),
1064        "usage": Value::Object(usage),
1065    })
1066}
1067
1068fn question_json(question: &QuestionReport) -> Value {
1069    match question {
1070        QuestionReport::Noul {
1071            cases,
1072            brier,
1073            accuracy,
1074            best,
1075            sweep,
1076            ..
1077        } => json!({
1078            "kind": question.kind(),
1079            "cases": cases,
1080            "brier": number(*brier),
1081            "accuracy": number(*accuracy),
1082            "best": {"threshold": number(best.threshold), "f1": number(best.f1)},
1083            "sweep": sweep.iter().map(|row| json!({
1084                "threshold": number(row.threshold),
1085                "tp": row.tp,
1086                "fp": row.fp,
1087                "fn": row.r#fn,
1088                "tn": row.tn,
1089                "accuracy": number(row.accuracy),
1090                "precision": maybe(row.precision),
1091                "recall": maybe(row.recall),
1092                "f1": number(row.f1),
1093            })).collect::<Vec<_>>(),
1094        }),
1095        QuestionReport::Choice {
1096            cases,
1097            accuracy,
1098            labels,
1099            confusion,
1100            gate,
1101            ..
1102        } => json!({
1103            "kind": question.kind(),
1104            "cases": cases,
1105            "accuracy": number(*accuracy),
1106            "labels": labels,
1107            "confusion": confusion,
1108            "gate": gate.iter().map(gate_json).collect::<Vec<_>>(),
1109        }),
1110        QuestionReport::Score {
1111            cases,
1112            exact,
1113            within_one,
1114            mae,
1115            gate,
1116            ..
1117        } => json!({
1118            "kind": question.kind(),
1119            "cases": cases,
1120            "exact": number(*exact),
1121            "withinOne": number(*within_one),
1122            "mae": number(*mae),
1123            "gate": gate.iter().map(gate_json).collect::<Vec<_>>(),
1124        }),
1125    }
1126}
1127
1128fn gate_json(row: &GateRow) -> Value {
1129    json!({
1130        "confidence": number(row.confidence),
1131        "coverage": number(row.coverage),
1132        "accuracy": maybe(row.accuracy),
1133    })
1134}
1135
1136/// A number written the way `JSON.stringify` writes it: a whole float loses its `.0`, so the two
1137/// ports' JSON reports can be compared byte for byte the way their tables can.
1138fn number(x: f64) -> Value {
1139    if x.fract() == 0.0 && x.abs() < 9e15 {
1140        return json!(x as i64);
1141    }
1142    json!(x)
1143}
1144
1145/// The same, for a rate that was never defined: `null`, so the key is always there.
1146fn maybe(x: Option<f64>) -> Value {
1147    x.map_or(Value::Null, number)
1148}
1149
1150/// A rate that was never defined is a dot, not a zero: nothing was measured.
1151fn rate(n: Option<f64>) -> String {
1152    match n {
1153        Some(n) => two(n),
1154        None => "·".to_owned(),
1155    }
1156}
1157
1158fn plural(n: usize) -> &'static str {
1159    if n == 1 { "" } else { "s" }
1160}
1161
1162fn pad_end(text: &str, width: usize) -> String {
1163    let length = text.chars().count();
1164    if length >= width {
1165        text.to_owned()
1166    } else {
1167        format!("{text}{}", " ".repeat(width - length))
1168    }
1169}
1170
1171/// Styled lines as the plain text a pipe wants.
1172pub fn report_text(report: &Report) -> String {
1173    let mut out = String::new();
1174    for line in report_lines(report) {
1175        for span in &line.spans {
1176            out.push_str(span.content.as_ref());
1177        }
1178        out.push('\n');
1179    }
1180    out
1181}