1//! Answers after verification: every value in range, every label or
2//! level one that was asked, distributions in the question's order.
3//! `confidence` is kept exactly as the vendor returned it.
4
5use serde::{Deserialize, Serialize};
6
7/// Dollars per million INPUT tokens; output tokens are free. The vendor's
8/// published price (jevsnes `research/jev-api-digest.md`: $42 per billion,
9/// $0.042 per million), and postjevsql's default for `jev.price_per_mtok`.
10pub const DOLLARS_PER_MTOK: f64 = 0.042;
11
12#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
13pub struct NoulAnswer {
14    /// Probability that the answer is yes.
15    pub noul: f64,
16}
17
18#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
19pub struct ChoiceAnswer {
20    pub choice: String,
21    pub confidence: f64,
22    /// Every option with its probability, in the order they were asked.
23    pub probabilities: Vec<(String, f64)>,
24}
25
26#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
27pub struct ScoreAnswer {
28    /// The probability-weighted level, from 0 to levels − 1.
29    pub score: f64,
30    pub confidence: f64,
31    /// Probability of each level, lowest first.
32    pub probabilities: Vec<f64>,
33    /// The vendor's description of each level, lowest first.
34    pub legend: Vec<String>,
35}
36
37impl ScoreAnswer {
38    /// `score` on 0 to 1, so rubrics of different lengths compare.
39    pub fn normalized(&self) -> f64 {
40        self.score / (self.probabilities.len() - 1) as f64
41    }
42}
43
44#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
45pub struct Usage {
46    pub input_tokens: u64,
47    pub output_tokens: u64,
48}
49
50impl Usage {
51    /// What this usage cost at the published price: input tokens only.
52    pub fn dollars(self) -> f64 {
53        self.input_tokens as f64 / 1e6 * DOLLARS_PER_MTOK
54    }
55}
56
57impl std::ops::AddAssign for Usage {
58    fn add_assign(&mut self, other: Self) {
59        self.input_tokens += other.input_tokens;
60        self.output_tokens += other.output_tokens;
61    }
62}
63
64#[derive(Clone, Debug, PartialEq)]
65pub enum Answer {
66    Noul(NoulAnswer),
67    Choice(ChoiceAnswer),
68    Score(ScoreAnswer),
69}
70
71/// Which answer a `Key<A>` reads. Sealed: only the three answer types.
72pub trait FromAnswer: sealed::Sealed + Sized {
73    #[doc(hidden)]
74    fn from_answer(answer: &Answer) -> Option<&Self>;
75}
76
77mod sealed {
78    pub trait Sealed {}
79    impl Sealed for super::NoulAnswer {}
80    impl Sealed for super::ChoiceAnswer {}
81    impl Sealed for super::ScoreAnswer {}
82}
83
84impl FromAnswer for NoulAnswer {
85    fn from_answer(answer: &Answer) -> Option<&Self> {
86        match answer {
87            Answer::Noul(a) => Some(a),
88            _ => None,
89        }
90    }
91}
92
93impl FromAnswer for ChoiceAnswer {
94    fn from_answer(answer: &Answer) -> Option<&Self> {
95        match answer {
96            Answer::Choice(a) => Some(a),
97            _ => None,
98        }
99    }
100}
101
102impl FromAnswer for ScoreAnswer {
103    fn from_answer(answer: &Answer) -> Option<&Self> {
104        match answer {
105            Answer::Score(a) => Some(a),
106            _ => None,
107        }
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn a_million_input_tokens_costs_the_published_price() {
117        let usage = Usage { input_tokens: 1_000_000, output_tokens: 999 };
118        assert!((usage.dollars() - DOLLARS_PER_MTOK).abs() < 1e-12);
119    }
120
121    #[test]
122    fn an_answer_round_trips_through_a_record() {
123        let answer = ChoiceAnswer {
124            choice: "b".into(),
125            confidence: 0.4,
126            probabilities: vec![("b".into(), 0.6), ("a".into(), 0.4)],
127        };
128        let line = serde_json::to_string(&answer).unwrap();
129        assert_eq!(serde_json::from_str::<ChoiceAnswer>(&line).unwrap(), answer);
130    }
131}