cortexkit_role_classifier/
identity.rs1use serde::{Deserialize, Serialize};
13use sha2::{Digest, Sha256};
14
15use crate::request::ClassifyRequest;
16
17#[derive(Clone, Debug, PartialEq, Eq, Deserialize, Serialize)]
19#[non_exhaustive]
20pub struct BatchIdentity {
21 pub model: String,
23 pub questions: String,
25 pub items: Vec<String>,
27}
28
29#[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
42pub 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
47pub 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 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 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 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}