Skip to main content

cortexkit_role_classifier/
identity.rs

1//! The batch body identity: what decides whether a re-sent `batch_id` is a
2//! replay or `batch_id_reuse`.
3//!
4//! The identity is the model, the questions, and each item, the questions
5//! and every item hashed as RFC 8785 (JCS) canonical JSON. Canonicalizing
6//! first means the key order of an object `state`, whitespace, and number
7//! spellings JSON treats as equal (`1.0` and `1`, `1e30` and `1E+30`) can't
8//! turn a retry into a reuse. `max_cost_usd` is not part of it: a re-send
9//! may carry another ceiling (see
10//! [`effective_ceiling`](crate::request::effective_ceiling)).
11
12use serde::{Deserialize, Serialize};
13use sha2::{Digest, Sha256};
14
15use crate::request::ClassifyRequest;
16
17/// A batch's body identity, as a runner records it at first admission.
18#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
19#[non_exhaustive]
20pub struct BatchIdentity {
21    /// The catalog id, compared as is.
22    pub model: String,
23    /// SHA-256 of the JCS of `questions`, lowercase hex.
24    pub questions: String,
25    /// SHA-256 of the JCS of each item, lowercase hex, in request order.
26    pub items: Vec<String>,
27}
28
29/// Why an identity could not be computed: the request holds a value JCS
30/// cannot encode.
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct IdentityError(pub String);
33
34impl std::fmt::Display for IdentityError {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(f, "request cannot be canonicalized: {}", self.0)
37    }
38}
39
40impl std::error::Error for IdentityError {}
41
42/// The RFC 8785 (JCS) canonical JSON bytes of `value`.
43pub fn canonical_json<T: Serialize + ?Sized>(value: &T) -> Result<Vec<u8>, IdentityError> {
44    serde_jcs::to_vec(value).map_err(|e| IdentityError(e.to_string()))
45}
46
47/// SHA-256 of the canonical JSON of `value`, as 64 lowercase hex characters.
48pub fn digest<T: Serialize + ?Sized>(value: &T) -> Result<String, IdentityError> {
49    let digest = Sha256::digest(canonical_json(value)?);
50    Ok(digest.iter().map(|byte| format!("{byte:02x}")).collect())
51}
52
53impl BatchIdentity {
54    /// The identity of `request`.
55    pub fn of(request: &ClassifyRequest) -> Result<Self, IdentityError> {
56        Ok(Self {
57            model: request.model.clone(),
58            questions: digest(&request.questions)?,
59            items: request.items.iter().map(digest).collect::<Result<_, _>>()?,
60        })
61    }
62
63    /// The first field in which `sent` differs from this recorded identity,
64    /// as `batch_id_reuse` names it: `model`, `questions`, `items` (a
65    /// different item count) or `items[i]`. `None` when they are the same
66    /// body, which makes the re-send a replay.
67    pub fn first_difference(&self, sent: &Self) -> Option<String> {
68        if self.model != sent.model {
69            return Some("model".into());
70        }
71        if self.questions != sent.questions {
72            return Some("questions".into());
73        }
74        if self.items.len() != sent.items.len() {
75            return Some("items".into());
76        }
77        self.items
78            .iter()
79            .zip(&sent.items)
80            .position(|(recorded, sent)| recorded != sent)
81            .map(|i| format!("items[{i}]"))
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use serde_json::{json, Number};
88
89    use super::*;
90    use crate::{
91        question::Question,
92        request::{parse_request, Item},
93    };
94
95    fn request(states: Vec<serde_json::Value>) -> ClassifyRequest {
96        ClassifyRequest::new(
97            "b",
98            "cloudflare-workers-ai/@cf/cloudflare/clef",
99            [("q".to_owned(), Question::noul("i"))].into(),
100            states.into_iter().map(Item::new).collect(),
101        )
102    }
103
104    #[test]
105    fn key_order_and_number_spelling_do_not_change_the_identity() {
106        let first = parse_request(
107            r#"{"batch_id":"b","model":"m","questions":{"q":{"type":"noul","instructions":"i"}},"items":[{"state":{"b":1.0,"a":1e30,"c":{"y":[2.50],"x":0}}}]}"#,
108        )
109        .unwrap();
110        let second = parse_request(
111            r#"{"batch_id":"b","model":"m","questions":{"q":{"type":"noul","instructions":"i"}},"items":[{"state":{"c":{"x":0,"y":[2.5]},"a":1E+30,"b":1}}]}"#,
112        )
113        .unwrap();
114        assert_eq!(
115            String::from_utf8(canonical_json(&first.items[0].state).unwrap()).unwrap(),
116            r#"{"a":1e+30,"b":1,"c":{"x":0,"y":[2.5]}}"#
117        );
118        assert_eq!(
119            BatchIdentity::of(&first).unwrap(),
120            BatchIdentity::of(&second).unwrap()
121        );
122    }
123
124    #[test]
125    fn plain_json_would_have_called_the_same_state_different() {
126        // The reason for canonicalizing: serde_json writes the float `1.0`
127        // and the integer `1` differently, so hashing its text would refuse
128        // a faithful retry as reuse.
129        let float: serde_json::Value = serde_json::from_str(r#"{"a":1.0}"#).unwrap();
130        let int: serde_json::Value = serde_json::from_str(r#"{"a":1}"#).unwrap();
131        assert_ne!(
132            serde_json::to_string(&float).unwrap(),
133            serde_json::to_string(&int).unwrap()
134        );
135        assert_eq!(digest(&float).unwrap(), digest(&int).unwrap());
136    }
137
138    #[test]
139    fn the_first_difference_is_named() {
140        let recorded = BatchIdentity::of(&request(vec![json!("a"), json!("b")])).unwrap();
141        let same = BatchIdentity::of(&request(vec![json!("a"), json!("b")])).unwrap();
142        assert_eq!(recorded.first_difference(&same), None);
143        let changed = BatchIdentity::of(&request(vec![json!("a"), json!("c")])).unwrap();
144        assert_eq!(recorded.first_difference(&changed), Some("items[1]".into()));
145        let longer = BatchIdentity::of(&request(vec![json!("a"), json!("b"), json!("c")])).unwrap();
146        assert_eq!(recorded.first_difference(&longer), Some("items".into()));
147        let mut other_questions = request(vec![json!("a"), json!("b")]);
148        other_questions
149            .questions
150            .insert("r".into(), Question::noul("j"));
151        assert_eq!(
152            recorded.first_difference(&BatchIdentity::of(&other_questions).unwrap()),
153            Some("questions".into())
154        );
155        let mut other_model = request(vec![json!("a"), json!("b")]);
156        other_model.model = "typesafe/jev-1.13.0".into();
157        assert_eq!(
158            recorded.first_difference(&BatchIdentity::of(&other_model).unwrap()),
159            Some("model".into())
160        );
161    }
162
163    #[test]
164    fn max_cost_usd_is_not_part_of_the_identity() {
165        let plain = request(vec![json!("a")]);
166        let capped = request(vec![json!("a")]).with_max_cost_usd(Number::from_f64(0.01).unwrap());
167        assert_eq!(
168            BatchIdentity::of(&plain).unwrap(),
169            BatchIdentity::of(&capped).unwrap()
170        );
171    }
172}