cortexkit_role_classifier/
reply.rs

1//! `classify.run` replies and the provider's answer objects.
2//!
3//! After admission the batch never fails as a whole: every item gets
4//! `answers` or `error`, at its position in the request.
5
6use indexmap::IndexMap;
7use serde::{Deserialize, Serialize};
8use serde_json::{Number, Value};
9
10use crate::{errors::ItemError, question::QuestionType, request::ClassifyRequest};
11
12/// One answer, the provider's answer object as returned. Which members are
13/// present depends on the question type:
14///
15/// - `noul`: `{type, noul}`. `noul` is the probability that the answer is
16///   yes, from 0 (no) to 1 (yes).
17/// - `choice`: `{type, choice, probabilities, confidence}`.
18/// - `score`: `{type, score, legend, probabilities, confidence}`. `score`
19///   is fractional, from 0 to levels - 1, and can land between levels;
20///   `legend` and `probabilities` are keyed by the level index as a string.
21///
22/// Numbers are [`Number`]s, so an integer stays an integer, and a member
23/// the provider adds that this crate does not name is kept in
24/// [`Answer::extra`]. A reply decodes open: an answer of a type this crate
25/// does not name keeps its type string, and every member is kept as given.
26/// The runner never rounds, rescales or renames them.
27#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
28#[non_exhaustive]
29pub struct Answer {
30    #[serde(rename = "type")]
31    pub kind: QuestionType,
32    #[serde(default, skip_serializing_if = "Option::is_none")]
33    pub noul: Option<Number>,
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub choice: Option<String>,
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub score: Option<Number>,
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub legend: Option<IndexMap<String, Value>>,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub probabilities: Option<IndexMap<String, Number>>,
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub confidence: Option<Number>,
44    /// Members the provider returned that this crate does not name, kept
45    /// in the order they came.
46    #[serde(flatten)]
47    pub extra: IndexMap<String, Value>,
48}
49
50impl Answer {
51    fn empty(kind: QuestionType) -> Self {
52        Self {
53            kind,
54            noul: None,
55            choice: None,
56            score: None,
57            legend: None,
58            probabilities: None,
59            confidence: None,
60            extra: IndexMap::new(),
61        }
62    }
63
64    /// A `noul` answer.
65    pub fn noul(noul: Number) -> Self {
66        Self {
67            noul: Some(noul),
68            ..Self::empty(QuestionType::Noul)
69        }
70    }
71
72    /// A `choice` answer.
73    pub fn choice(
74        choice: impl Into<String>,
75        probabilities: IndexMap<String, Number>,
76        confidence: Number,
77    ) -> Self {
78        Self {
79            choice: Some(choice.into()),
80            probabilities: Some(probabilities),
81            confidence: Some(confidence),
82            ..Self::empty(QuestionType::Choice)
83        }
84    }
85
86    /// A `score` answer.
87    pub fn score(
88        score: Number,
89        legend: IndexMap<String, Value>,
90        probabilities: IndexMap<String, Number>,
91        confidence: Number,
92    ) -> Self {
93        Self {
94            score: Some(score),
95            legend: Some(legend),
96            probabilities: Some(probabilities),
97            confidence: Some(confidence),
98            ..Self::empty(QuestionType::Score)
99        }
100    }
101}
102
103/// Token usage. A count the provider did not report is absent, never zero.
104#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize)]
105#[non_exhaustive]
106pub struct Usage {
107    #[serde(default, skip_serializing_if = "Option::is_none")]
108    pub input_tokens: Option<u64>,
109    #[serde(default, skip_serializing_if = "Option::is_none")]
110    pub output_tokens: Option<u64>,
111}
112
113impl Usage {
114    pub fn new() -> Self {
115        Self::default()
116    }
117
118    pub fn with_input_tokens(mut self, tokens: u64) -> Self {
119        self.input_tokens = Some(tokens);
120        self
121    }
122
123    pub fn with_output_tokens(mut self, tokens: u64) -> Self {
124        self.output_tokens = Some(tokens);
125        self
126    }
127
128    /// Whether neither count was reported.
129    pub fn is_empty(&self) -> bool {
130        self.input_tokens.is_none() && self.output_tokens.is_none()
131    }
132}
133
134/// The batch's usage: each count summed over the items that reported it,
135/// and absent when none did.
136pub fn sum_usage<'a>(items: impl IntoIterator<Item = &'a ItemResult>) -> Usage {
137    // Absent is not zero (`cortexkit-model-batch`, its `RULES.md` §10).
138    let add = cortexkit_model_batch::absent::add_reported;
139    items
140        .into_iter()
141        .filter_map(|item| item.usage)
142        .fold(Usage::default(), |total, usage| Usage {
143            input_tokens: add(total.input_tokens, usage.input_tokens),
144            output_tokens: add(total.output_tokens, usage.output_tokens),
145        })
146}
147
148/// One item's outcome: `answers` if it succeeded, `error` if it failed,
149/// never both.
150#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
151#[non_exhaustive]
152pub struct ItemResult {
153    /// The item's position in the request.
154    pub index: u64,
155    #[serde(default, skip_serializing_if = "Option::is_none")]
156    pub answers: Option<IndexMap<String, Answer>>,
157    #[serde(default, skip_serializing_if = "Option::is_none")]
158    pub error: Option<ItemError>,
159    /// What the provider reported for this item; absent when it reported
160    /// nothing.
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub usage: Option<Usage>,
163}
164
165impl ItemResult {
166    pub fn answered(index: u64, answers: IndexMap<String, Answer>) -> Self {
167        Self {
168            index,
169            answers: Some(answers),
170            error: None,
171            usage: None,
172        }
173    }
174
175    pub fn failed(index: u64, error: ItemError) -> Self {
176        Self {
177            index,
178            answers: None,
179            error: Some(error),
180            usage: None,
181        }
182    }
183
184    pub fn with_usage(mut self, usage: Usage) -> Self {
185        self.usage = Some(usage);
186        self
187    }
188}
189
190/// The `classify.run` reply.
191#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
192#[non_exhaustive]
193pub struct ClassifyReply {
194    pub batch_id: String,
195    pub model: String,
196    /// One per request item, in request order.
197    pub items: Vec<ItemResult>,
198    /// Summed over answered items; a count no item reported is absent.
199    #[serde(default)]
200    pub usage: Usage,
201    /// The batch's cumulative cost in USD, the same on every reply to the
202    /// batch, a replay included. Absent only when the model is unpriced.
203    #[serde(default, skip_serializing_if = "Option::is_none")]
204    pub cost_usd: Option<Number>,
205}
206
207impl ClassifyReply {
208    /// A reply whose usage is summed from `items`.
209    pub fn new(
210        batch_id: impl Into<String>,
211        model: impl Into<String>,
212        items: Vec<ItemResult>,
213    ) -> Self {
214        let usage = sum_usage(&items);
215        Self {
216            batch_id: batch_id.into(),
217            model: model.into(),
218            items,
219            usage,
220            cost_usd: None,
221        }
222    }
223
224    pub fn with_usage(mut self, usage: Usage) -> Self {
225        self.usage = usage;
226        self
227    }
228
229    pub fn with_cost_usd(mut self, cost_usd: Number) -> Self {
230        self.cost_usd = Some(cost_usd);
231        self
232    }
233}
234
235/// Check a reply against the request it answers: the same batch and model,
236/// one item per request item in request order, each with exactly one of
237/// `answers` and `error`, and every answered item answering exactly the
238/// request's questions with answers of the asked type. The error names the
239/// first problem.
240pub fn check_reply(reply: &ClassifyReply, request: &ClassifyRequest) -> Result<(), String> {
241    if reply.batch_id != request.batch_id {
242        return Err(format!(
243            "batch_id {} answers request {}",
244            reply.batch_id, request.batch_id
245        ));
246    }
247    if reply.model != request.model {
248        return Err(format!(
249            "model {} answers a request for {}",
250            reply.model, request.model
251        ));
252    }
253    if reply.items.len() != request.items.len() {
254        return Err(format!(
255            "{} items answer {} request items",
256            reply.items.len(),
257            request.items.len()
258        ));
259    }
260    for (position, item) in reply.items.iter().enumerate() {
261        if item.index != position as u64 {
262            return Err(format!(
263                "item at position {position} has index {}",
264                item.index
265            ));
266        }
267        match (&item.answers, &item.error) {
268            (Some(answers), None) => {
269                let asked: Vec<&String> = request.questions.keys().collect();
270                let got: Vec<&String> = answers.keys().collect();
271                let mut asked_sorted = asked.clone();
272                let mut got_sorted = got.clone();
273                asked_sorted.sort();
274                got_sorted.sort();
275                if asked_sorted != got_sorted {
276                    return Err(format!(
277                        "item {position} answers {got:?}, the request asked {asked:?}"
278                    ));
279                }
280                for (id, answer) in answers {
281                    let kind = &request.questions[id].kind;
282                    if &answer.kind != kind {
283                        return Err(format!(
284                            "item {position} answers {id} as {}, asked as {kind}",
285                            answer.kind
286                        ));
287                    }
288                }
289            }
290            (None, Some(_)) => {}
291            (Some(_), Some(_)) => {
292                return Err(format!("item {position} has both answers and error"))
293            }
294            (None, None) => return Err(format!("item {position} has neither answers nor error")),
295        }
296    }
297    Ok(())
298}
299
300#[cfg(test)]
301mod tests {
302    use serde_json::json;
303
304    use super::*;
305    use crate::{errors::ItemError, question::Question, request::Item};
306
307    #[test]
308    fn unreported_usage_stays_absent_in_the_sum() {
309        let items = vec![
310            ItemResult::answered(0, IndexMap::new()),
311            ItemResult::answered(1, IndexMap::new()).with_usage(Usage::new().with_input_tokens(5)),
312            ItemResult::failed(2, ItemError::for_code("invalid_item", "m")),
313        ];
314        let usage = sum_usage(&items);
315        assert_eq!(usage.input_tokens, Some(5));
316        assert_eq!(usage.output_tokens, None);
317        assert!(sum_usage(&items[..1]).is_empty());
318        let reply = ClassifyReply::new("b", "m", items[..1].to_vec());
319        assert_eq!(
320            serde_json::to_string(&reply).unwrap(),
321            r#"{"batch_id":"b","model":"m","items":[{"index":0,"answers":{}}],"usage":{}}"#
322        );
323    }
324
325    #[test]
326    fn answers_keep_integers_order_and_unknown_members() {
327        let text = r#"{"type":"choice","choice":"b","probabilities":{"b":1,"a":0.0},"confidence":0.5,"rationale":"x"}"#;
328        let answer: Answer = serde_json::from_str(text).unwrap();
329        assert_eq!(answer.extra["rationale"], json!("x"));
330        assert_eq!(serde_json::to_string(&answer).unwrap(), text);
331    }
332
333    #[test]
334    fn a_reply_out_of_order_or_with_both_members_is_refused() {
335        let request = ClassifyRequest::new(
336            "b",
337            "m",
338            [("q".to_owned(), Question::noul("i"))].into(),
339            vec![Item::new("a"), Item::new("b")],
340        );
341        let answers: IndexMap<String, Answer> =
342            [("q".to_owned(), Answer::noul(Number::from_f64(0.5).unwrap()))].into();
343        let good = ClassifyReply::new(
344            "b",
345            "m",
346            vec![
347                ItemResult::answered(0, answers.clone()),
348                ItemResult::failed(1, ItemError::for_code("rate_limited", "m")),
349            ],
350        );
351        assert_eq!(check_reply(&good, &request), Ok(()));
352        let mut swapped = good.clone();
353        swapped.items.swap(0, 1);
354        assert!(check_reply(&swapped, &request)
355            .unwrap_err()
356            .contains("index"));
357        let mut both = good;
358        both.items[1].answers = Some(answers);
359        assert!(check_reply(&both, &request).unwrap_err().contains("both"));
360    }
361}
cortexkit_role_classifier/reply.rs

cortexkit_role_classifier/
reply.rs
