1use std::collections::{BTreeMap, HashMap};
19
20use async_trait::async_trait;
21use serde::{Deserialize, Serialize};
22
23use crate::{AgentLoopError, Result};
24
25#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
27pub enum ClassificationQuestion {
28 Noul {
30 instructions: String,
32 yes: Option<String>,
34 no: Option<String>,
36 },
37 Choice {
39 instructions: String,
41 options: Vec<(String, Option<String>)>,
43 },
44 Score {
46 instructions: String,
48 levels: Vec<String>,
50 },
51}
52
53impl ClassificationQuestion {
54 pub fn noul(instructions: impl Into<String>) -> Self {
56 Self::Noul {
57 instructions: instructions.into(),
58 yes: None,
59 no: None,
60 }
61 }
62
63 pub fn score<L: Into<String>>(
65 instructions: impl Into<String>,
66 levels: impl IntoIterator<Item = L>,
67 ) -> Self {
68 Self::Score {
69 instructions: instructions.into(),
70 levels: levels.into_iter().map(Into::into).collect(),
71 }
72 }
73
74 pub fn kind(&self) -> &'static str {
76 match self {
77 Self::Noul { .. } => "noul",
78 Self::Choice { .. } => "choice",
79 Self::Score { .. } => "score",
80 }
81 }
82}
83
84#[derive(Debug, Clone, Default)]
86pub struct ClassificationRequest {
87 pub state: serde_json::Value,
89 pub questions: Vec<(String, ClassificationQuestion)>,
92 pub metadata: HashMap<String, String>,
94}
95
96impl ClassificationRequest {
97 pub fn new(state: impl Into<serde_json::Value>) -> Self {
99 Self {
100 state: state.into(),
101 questions: Vec::new(),
102 metadata: HashMap::new(),
103 }
104 }
105
106 pub fn ask(mut self, id: impl Into<String>, question: ClassificationQuestion) -> Self {
108 self.questions.push((id.into(), question));
109 self
110 }
111
112 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
114 self.metadata.insert(key.into(), value.into());
115 self
116 }
117
118 pub fn len(&self) -> usize {
120 self.questions.len()
121 }
122
123 pub fn is_empty(&self) -> bool {
125 self.questions.is_empty()
126 }
127}
128
129#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
131pub enum ClassificationAnswer {
132 Noul {
137 probability: f64,
139 },
140 Choice {
142 selected: String,
144 probabilities: BTreeMap<String, f64>,
146 confidence: f64,
148 },
149 Score {
151 score: f64,
153 probabilities: BTreeMap<usize, f64>,
155 confidence: f64,
157 },
158}
159
160impl ClassificationAnswer {
161 pub fn kind(&self) -> &'static str {
163 match self {
164 Self::Noul { .. } => "noul",
165 Self::Choice { .. } => "choice",
166 Self::Score { .. } => "score",
167 }
168 }
169
170 pub fn probability_yes(&self) -> Option<f64> {
172 match self {
173 Self::Noul { probability } => Some(*probability),
174 _ => None,
175 }
176 }
177
178 pub fn confidence(&self) -> Option<f64> {
180 match self {
181 Self::Noul { .. } => None,
182 Self::Choice { confidence, .. } | Self::Score { confidence, .. } => Some(*confidence),
183 }
184 }
185
186 pub fn probability_at_or_above(&self, level: usize) -> Option<f64> {
192 match self {
193 Self::Score { probabilities, .. } => Some(
194 probabilities
195 .iter()
196 .filter(|(index, _)| **index >= level)
197 .map(|(_, probability)| probability)
198 .sum(),
199 ),
200 _ => None,
201 }
202 }
203}
204
205#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
207pub struct ClassificationUsage {
208 pub input_tokens: u64,
210 pub output_tokens: u64,
212}
213
214#[derive(Debug, Clone, Default, PartialEq)]
216pub struct ClassificationOutcome {
217 pub model: String,
219 pub answers: BTreeMap<String, ClassificationAnswer>,
221 pub usage: ClassificationUsage,
223}
224
225impl ClassificationOutcome {
226 pub fn get(&self, id: &str) -> Option<&ClassificationAnswer> {
228 self.answers.get(id)
229 }
230}
231
232#[async_trait]
234pub trait ClassifierService: Send + Sync {
235 fn is_configured(&self) -> bool;
237
238 async fn evaluate(&self, request: ClassificationRequest) -> Result<ClassificationOutcome>;
240
241 fn name(&self) -> &'static str {
243 "ClassifierService"
244 }
245}
246
247#[derive(Debug, Clone, Default)]
249pub struct DisabledClassifierService;
250
251#[async_trait]
252impl ClassifierService for DisabledClassifierService {
253 fn is_configured(&self) -> bool {
254 false
255 }
256
257 async fn evaluate(&self, _request: ClassificationRequest) -> Result<ClassificationOutcome> {
258 Err(AgentLoopError::llm(
259 "classifier is disabled (no UTILITY_TYPESAFE_API_KEY configured)",
260 ))
261 }
262
263 fn name(&self) -> &'static str {
264 "DisabledClassifierService"
265 }
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271
272 #[tokio::test]
273 async fn disabled_service_reports_itself_and_rejects_requests() {
274 let service = DisabledClassifierService;
275 assert!(!service.is_configured());
276 let error = service
277 .evaluate(
278 ClassificationRequest::new("anything")
279 .ask("q", ClassificationQuestion::noul("Yes?")),
280 )
281 .await
282 .unwrap_err();
283 assert!(
284 error.to_string().contains("UTILITY_TYPESAFE_API_KEY"),
285 "a disabled service must name the switch that enables it: {error}"
286 );
287 }
288
289 #[test]
290 fn requests_keep_question_order_and_carry_metadata() {
291 let request = ClassificationRequest::new(serde_json::json!({"text": "hi"}))
292 .ask("second", ClassificationQuestion::noul("b"))
293 .ask("first", ClassificationQuestion::score("a", ["low", "high"]))
294 .with_metadata("purpose", "guardrails");
295 assert_eq!(request.len(), 2);
296 assert_eq!(request.questions[0].0, "second");
297 assert_eq!(request.questions[1].1.kind(), "score");
298 assert_eq!(request.metadata["purpose"], "guardrails");
299 assert_eq!(request.state["text"], "hi");
300 }
301
302 #[test]
303 fn score_tail_mass_is_read_from_the_distribution_not_the_mean() {
304 let answer = ClassificationAnswer::Score {
306 score: 0.6,
307 probabilities: BTreeMap::from([(0, 0.7), (1, 0.0), (2, 0.3)]),
308 confidence: 0.4,
309 };
310 assert_eq!(answer.probability_at_or_above(2), Some(0.3));
311 assert_eq!(answer.probability_at_or_above(0), Some(1.0));
312 assert_eq!(answer.probability_yes(), None);
313 assert_eq!(answer.confidence(), Some(0.4));
314 }
315
316 #[test]
317 fn noul_answers_have_no_separate_confidence() {
318 let answer = ClassificationAnswer::Noul { probability: 0.92 };
319 assert_eq!(answer.probability_yes(), Some(0.92));
320 assert_eq!(answer.confidence(), None);
321 assert_eq!(answer.probability_at_or_above(1), None);
322 }
323}