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    let add = |total: Option<u64>, count: Option<u64>| match (total, count) {
138        (Some(total), Some(count)) => Some(total + count),
139        (total, None) => total,
140        (None, count) => count,
141    };
142    items
143        .into_iter()
144        .filter_map(|item| item.usage)
145        .fold(Usage::default(), |total, usage| Usage {
146            input_tokens: add(total.input_tokens, usage.input_tokens),
147            output_tokens: add(total.output_tokens, usage.output_tokens),
148        })
149}
150
151/// One item's outcome: `answers` if it succeeded, `error` if it failed,
152/// never both.
153#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
154#[non_exhaustive]
155pub struct ItemResult {
156    /// The item's position in the request.
157    pub index: u64,
158    #[serde(default, skip_serializing_if = "Option::is_none")]
159    pub answers: Option<IndexMap<String, Answer>>,
160    #[serde(default, skip_serializing_if = "Option::is_none")]
161    pub error: Option<ItemError>,
162    /// What the provider reported for this item; absent when it reported
163    /// nothing.
164    #[serde(default, skip_serializing_if = "Option::is_none")]
165    pub usage: Option<Usage>,
166}
167
168impl ItemResult {
169    pub fn answered(index: u64, answers: IndexMap<String, Answer>) -> Self {
170        Self {
171            index,
172            answers: Some(answers),
173            error: None,
174            usage: None,
175        }
176    }
177
178    pub fn failed(index: u64, error: ItemError) -> Self {
179        Self {
180            index,
181            answers: None,
182            error: Some(error),
183            usage: None,
184        }
185    }
186
187    pub fn with_usage(mut self, usage: Usage) -> Self {
188        self.usage = Some(usage);
189        self
190    }
191}
192
193/// The `classify.run` reply.
194#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
195#[non_exhaustive]
196pub struct ClassifyReply {
197    pub batch_id: String,
198    pub model: String,
199    /// One per request item, in request order.
200    pub items: Vec<ItemResult>,
201    /// Summed over answered items; a count no item reported is absent.
202    #[serde(default)]
203    pub usage: Usage,
204    /// The batch's cumulative cost in USD, the same on every reply to the
205    /// batch, a replay included. Absent only when the model is unpriced.
206    #[serde(default, skip_serializing_if = "Option::is_none")]
207    pub cost_usd: Option<Number>,
208}
209
210impl ClassifyReply {
211    /// A reply whose usage is summed from `items`.
212    pub fn new(
213        batch_id: impl Into<String>,
214        model: impl Into<String>,
215        items: Vec<ItemResult>,
216    ) -> Self {
217        let usage = sum_usage(&items);
218        Self {
219            batch_id: batch_id.into(),
220            model: model.into(),
221            items,
222            usage,
223            cost_usd: None,
224        }
225    }
226
227    pub fn with_usage(mut self, usage: Usage) -> Self {
228        self.usage = usage;
229        self
230    }
231
232    pub fn with_cost_usd(mut self, cost_usd: Number) -> Self {
233        self.cost_usd = Some(cost_usd);
234        self
235    }
236}
237
238/// Check a reply against the request it answers: the same batch and model,
239/// one item per request item in request order, each with exactly one of
240/// `answers` and `error`, and every answered item answering exactly the
241/// request's questions with answers of the asked type. The error names the
242/// first problem.
243pub fn check_reply(reply: &ClassifyReply, request: &ClassifyRequest) -> Result<(), String> {
244    if reply.batch_id != request.batch_id {
245        return Err(format!(
246            "batch_id {} answers request {}",
247            reply.batch_id, request.batch_id
248        ));
249    }
250    if reply.model != request.model {
251        return Err(format!(
252            "model {} answers a request for {}",
253            reply.model, request.model
254        ));
255    }
256    if reply.items.len() != request.items.len() {
257        return Err(format!(
258            "{} items answer {} request items",
259            reply.items.len(),
260            request.items.len()
261        ));
262    }
263    for (position, item) in reply.items.iter().enumerate() {
264        if item.index != position as u64 {
265            return Err(format!(
266                "item at position {position} has index {}",
267                item.index
268            ));
269        }
270        match (&item.answers, &item.error) {
271            (Some(answers), None) => {
272                let asked: Vec<&String> = request.questions.keys().collect();
273                let got: Vec<&String> = answers.keys().collect();
274                let mut asked_sorted = asked.clone();
275                let mut got_sorted = got.clone();
276                asked_sorted.sort();
277                got_sorted.sort();
278                if asked_sorted != got_sorted {
279                    return Err(format!(
280                        "item {position} answers {got:?}, the request asked {asked:?}"
281                    ));
282                }
283                for (id, answer) in answers {
284                    let kind = &request.questions[id].kind;
285                    if &answer.kind != kind {
286                        return Err(format!(
287                            "item {position} answers {id} as {}, asked as {kind}",
288                            answer.kind
289                        ));
290                    }
291                }
292            }
293            (None, Some(_)) => {}
294            (Some(_), Some(_)) => {
295                return Err(format!("item {position} has both answers and error"))
296            }
297            (None, None) => return Err(format!("item {position} has neither answers nor error")),
298        }
299    }
300    Ok(())
301}
302
303#[cfg(test)]
304mod tests {
305    use serde_json::json;
306
307    use super::*;
308    use crate::{errors::ItemError, question::Question, request::Item};
309
310    #[test]
311    fn unreported_usage_stays_absent_in_the_sum() {
312        let items = vec![
313            ItemResult::answered(0, IndexMap::new()),
314            ItemResult::answered(1, IndexMap::new()).with_usage(Usage::new().with_input_tokens(5)),
315            ItemResult::failed(2, ItemError::for_code("invalid_item", "m")),
316        ];
317        let usage = sum_usage(&items);
318        assert_eq!(usage.input_tokens, Some(5));
319        assert_eq!(usage.output_tokens, None);
320        assert!(sum_usage(&items[..1]).is_empty());
321        let reply = ClassifyReply::new("b", "m", items[..1].to_vec());
322        assert_eq!(
323            serde_json::to_string(&reply).unwrap(),
324            r#"{"batch_id":"b","model":"m","items":[{"index":0,"answers":{}}],"usage":{}}"#
325        );
326    }
327
328    #[test]
329    fn answers_keep_integers_order_and_unknown_members() {
330        let text = r#"{"type":"choice","choice":"b","probabilities":{"b":1,"a":0.0},"confidence":0.5,"rationale":"x"}"#;
331        let answer: Answer = serde_json::from_str(text).unwrap();
332        assert_eq!(answer.extra["rationale"], json!("x"));
333        assert_eq!(serde_json::to_string(&answer).unwrap(), text);
334    }
335
336    #[test]
337    fn a_reply_out_of_order_or_with_both_members_is_refused() {
338        let request = ClassifyRequest::new(
339            "b",
340            "m",
341            [("q".to_owned(), Question::noul("i"))].into(),
342            vec![Item::new("a"), Item::new("b")],
343        );
344        let answers: IndexMap<String, Answer> =
345            [("q".to_owned(), Answer::noul(Number::from_f64(0.5).unwrap()))].into();
346        let good = ClassifyReply::new(
347            "b",
348            "m",
349            vec![
350                ItemResult::answered(0, answers.clone()),
351                ItemResult::failed(1, ItemError::for_code("rate_limited", "m")),
352            ],
353        );
354        assert_eq!(check_reply(&good, &request), Ok(()));
355        let mut swapped = good.clone();
356        swapped.items.swap(0, 1);
357        assert!(check_reply(&swapped, &request)
358            .unwrap_err()
359            .contains("index"));
360        let mut both = good;
361        both.items[1].answers = Some(answers);
362        assert!(check_reply(&both, &request).unwrap_err().contains("both"));
363    }
364}
cortexkit_role_classifier/reply.rs

cortexkit_role_classifier/
reply.rs
