cratefield_adapter_classifier_llm/
lib.rs1#![doc = include_str!("../README.md")]
30#![forbid(unsafe_code)]
31
32use cratefield_core::{
33 Answer, AnswerValue, Calibration, Classifier, ClassifierError, ClassifierProfile, Completion,
34 DEFAULT_MAX_STATE_CHARS, ModelTier, Prompt, Question, TextModel, TextModelError,
35 validate_questions,
36};
37use serde_json::{Map, Value};
38use std::collections::BTreeMap;
39use std::fmt::Write as _;
40use std::sync::Arc;
41
42const DEFAULT_TIER: ModelTier = ModelTier::Strong;
48
49const SYSTEM_PROMPT: &str = "You are a classification backend. You are given one state and a set of questions about it. For every question choose exactly one of the labels the question offers, and give a probability for every label the question offers; the probabilities of one question sum to 1. Answer with JSON conforming to the schema given, and with nothing else.";
51
52pub struct ClassifierLlm {
56 model: Arc<dyn TextModel>,
57 tier: ModelTier,
58 max_state_chars: usize,
59}
60
61impl ClassifierLlm {
62 #[must_use]
68 pub fn new(model: Arc<dyn TextModel>) -> Self {
69 Self {
70 model,
71 tier: DEFAULT_TIER,
72 max_state_chars: DEFAULT_MAX_STATE_CHARS,
73 }
74 }
75
76 #[must_use]
80 pub fn tier(mut self, tier: ModelTier) -> Self {
81 self.tier = tier;
82 self
83 }
84
85 #[must_use]
91 pub fn max_state_chars(mut self, max_state_chars: usize) -> Self {
92 self.max_state_chars = max_state_chars;
93 self
94 }
95}
96
97#[async_trait::async_trait]
98impl Classifier for ClassifierLlm {
99 fn profile(&self) -> ClassifierProfile {
100 ClassifierProfile::new(Calibration::LanguageModel, self.max_state_chars)
101 }
102
103 async fn ask(
104 &self,
105 state: &str,
106 questions: &BTreeMap<String, Question>,
107 ) -> Result<BTreeMap<String, Answer>, ClassifierError> {
108 validate_questions(questions)?;
112
113 let profile = self.profile();
114 let (state, truncated) = profile.truncate(state);
115 if truncated {
116 tracing::warn!(
119 limit = profile.max_state_chars,
120 "classifier state truncated to this adapter's state budget"
121 );
122 }
123
124 let prompt = Prompt::new(self.tier)
129 .system(SYSTEM_PROMPT)
130 .user(format!(
131 "STATE:\n{state}\nEND STATE.\n\n{}",
132 render_questions(questions)
133 ))
134 .json_schema(answer_schema(questions));
135
136 let completion = self.model.complete(&prompt).await.map_err(model_error)?;
137 let Some(answered) = answer_json(&completion) else {
138 return Err(ClassifierError::Transport(format!(
142 "the model answered {} chars, but neither the structured payload nor the text parsed as the answer object",
143 completion.text.len()
144 )));
145 };
146 answers_of(&answered, questions)
147 }
148}
149
150fn model_error(error: TextModelError) -> ClassifierError {
164 #[allow(clippy::match_same_arms)]
168 match error {
169 TextModelError::NotConfigured => ClassifierError::NotConfigured,
170 TextModelError::Rejected(message) => ClassifierError::Rejected(message),
171 TextModelError::Transient { retry_after } => ClassifierError::Transient { retry_after },
172 TextModelError::Transport(message) => ClassifierError::Transport(message),
173 TextModelError::Unsupported(_) => ClassifierError::NotConfigured,
174 _ => ClassifierError::NotConfigured,
177 }
178}
179
180fn render_questions(questions: &BTreeMap<String, Question>) -> String {
182 let mut rendered = String::from("QUESTIONS:\n");
183 for (index, (id, question)) in questions.iter().enumerate() {
184 let _ = writeln!(
185 rendered,
186 "{}. id {} ({}): {}",
187 index + 1,
188 id,
189 question.kind(),
190 question.instructions()
191 );
192 match question {
193 Question::Choice { criteria, .. } => {
194 rendered.push_str(" pick exactly one of:\n");
195 for (name, meaning) in criteria {
196 let _ = writeln!(rendered, " - {name}: {meaning}");
197 }
198 }
199 Question::Score { levels, .. } => {
200 rendered.push_str(" place on the ordered scale:\n");
201 for (name, meaning) in levels {
202 let _ = writeln!(rendered, " - {name}: {meaning}");
203 }
204 }
205 Question::Noul { .. } => {
206 rendered.push_str(" answer exactly true or false.\n");
207 }
208 }
209 }
210 rendered
211}
212
213fn answer_schema(questions: &BTreeMap<String, Question>) -> Value {
218 let mut properties = Map::new();
219 let mut required: Vec<Value> = Vec::new();
220 for (id, question) in questions {
221 let labels = question.labels();
222 let mut probability_properties = Map::new();
223 for label in &labels {
224 probability_properties
225 .insert((*label).to_owned(), serde_json::json!({ "type": "number" }));
226 }
227 properties.insert(
228 id.clone(),
229 serde_json::json!({
230 "type": "object",
231 "additionalProperties": false,
232 "required": ["value", "probabilities"],
233 "properties": {
234 "value": { "enum": labels },
235 "probabilities": {
236 "type": "object",
237 "additionalProperties": false,
238 "required": labels,
239 "properties": probability_properties,
240 },
241 },
242 }),
243 );
244 required.push(Value::String(id.clone()));
245 }
246 serde_json::json!({
247 "type": "object",
248 "additionalProperties": false,
249 "required": ["answers"],
250 "properties": {
251 "answers": {
252 "type": "object",
253 "additionalProperties": false,
254 "required": required,
255 "properties": properties,
256 },
257 },
258 })
259}
260
261fn answer_json(completion: &Completion) -> Option<Value> {
265 if let Some(json) = &completion.json {
266 return Some(json.clone());
267 }
268 serde_json::from_str(strip_code_fence(&completion.text)).ok()
269}
270
271fn strip_code_fence(text: &str) -> &str {
274 let trimmed = text.trim();
275 let Some(after_open) = trimmed.strip_prefix("```") else {
276 return trimmed;
277 };
278 let body = match after_open.split_once('\n') {
279 Some((_, body)) => body,
280 None => after_open,
281 };
282 match body.rfind("```") {
283 Some(end) => body[..end].trim(),
284 None => body.trim(),
285 }
286}
287
288fn answers_of(
294 answered: &Value,
295 questions: &BTreeMap<String, Question>,
296) -> Result<BTreeMap<String, Answer>, ClassifierError> {
297 let Some(answers) = answered.get("answers").and_then(Value::as_object) else {
298 return Err(ClassifierError::Rejected(
299 "the model's answer carries no `answers` object".to_owned(),
300 ));
301 };
302 for id in answers.keys() {
306 if !questions.contains_key(id) {
307 return Err(ClassifierError::Rejected(format!(
308 "the model answered question {id:?}, which was not asked"
309 )));
310 }
311 }
312 let mut out = BTreeMap::new();
313 for (id, question) in questions {
314 let Some(entry) = answers.get(id) else {
315 return Err(ClassifierError::Rejected(format!(
316 "the model answered no question {id:?}"
317 )));
318 };
319 out.insert(id.clone(), answer_for(question, entry)?);
320 }
321 Ok(out)
322}
323
324fn answer_for(question: &Question, entry: &Value) -> Result<Answer, ClassifierError> {
329 let object = entry.as_object().ok_or_else(|| {
330 ClassifierError::Rejected("the answer to a question is not an object".to_owned())
331 })?;
332 let chosen = object.get("value").and_then(Value::as_str).ok_or_else(|| {
333 ClassifierError::Rejected("the answer carries no chosen label as a string".to_owned())
334 })?;
335 let reported = object
336 .get("probabilities")
337 .and_then(Value::as_object)
338 .ok_or_else(|| {
339 ClassifierError::Rejected("the answer carries no `probabilities` object".to_owned())
340 })?;
341
342 let labels = question.labels();
343 let mut weights = BTreeMap::<String, f64>::new();
344 for (label, weight) in reported {
345 if !labels.contains(&label.as_str()) {
346 return Err(ClassifierError::Rejected(format!(
347 "the model gave a probability under {label:?}, a label the question never offered"
348 )));
349 }
350 let weight = weight.as_f64().ok_or_else(|| {
351 ClassifierError::Rejected(format!("the probability under {label:?} is not a number"))
352 })?;
353 if !weight.is_finite() {
357 return Err(ClassifierError::Rejected(format!(
358 "the probability under {label:?} is not finite"
359 )));
360 }
361 if weight < 0.0 {
362 return Err(ClassifierError::Rejected(format!(
363 "the probability under {label:?} is negative"
364 )));
365 }
366 weights.insert(label.clone(), weight);
367 }
368 for label in &labels {
369 if !weights.contains_key(*label) {
370 return Err(ClassifierError::Rejected(format!(
371 "the model gave no probability for the offered label {label:?}"
372 )));
373 }
374 }
375 let total: f64 = weights.values().sum();
376 if !(total.is_finite() && total > 0.0) {
377 return Err(ClassifierError::Rejected(
378 "the probabilities do not add up to a distribution".to_owned(),
379 ));
380 }
381 let mut probabilities = BTreeMap::<String, f32>::new();
382 for (label, weight) in &weights {
383 #[allow(clippy::cast_possible_truncation)]
387 let narrowed = (weight / total) as f32;
388 probabilities.insert(label.clone(), narrowed);
389 }
390 if !probabilities.contains_key(chosen) {
391 return Err(ClassifierError::Rejected(format!(
392 "the model chose {chosen:?}, a label the question never offered"
393 )));
394 }
395
396 match question {
397 Question::Choice { .. } => Ok(Answer::choice(chosen, probabilities)),
398 Question::Score { .. } => {
399 let parsed: f32 = chosen.parse().map_err(|_| {
400 ClassifierError::Rejected(format!(
401 "the score question chose {chosen:?}, which is not a number on its scale"
402 ))
403 })?;
404 if parsed.to_string() == chosen {
411 Ok(Answer::score(parsed, probabilities))
412 } else {
413 let confidence = probabilities[chosen];
414 Ok(Answer::new(
415 AnswerValue::Score(parsed),
416 probabilities,
417 confidence,
418 ))
419 }
420 }
421 Question::Noul { .. } => {
422 let value = chosen == "true";
425 Ok(Answer::noul(value, probabilities))
426 }
427 }
428}
429
430#[cfg(test)]
431mod tests {
432 use super::*;
433
434 #[test]
435 fn a_code_fence_is_stripped_and_bare_json_is_left_alone() {
436 let fenced = "```json\n{\"answers\": {}}\n```";
437 assert_eq!(strip_code_fence(fenced), "{\"answers\": {}}");
438 let bare_fence = "```\n{\"answers\": {}}\n```";
439 assert_eq!(strip_code_fence(bare_fence), "{\"answers\": {}}");
440 let padded = " \n{\"answers\": {}}\n ";
441 assert_eq!(strip_code_fence(padded), "{\"answers\": {}}");
442 assert_eq!(strip_code_fence("{\"answers\": {}}"), "{\"answers\": {}}");
443 assert_eq!(strip_code_fence("```json\n{\"a\": 1}"), "{\"a\": 1}");
446 }
447
448 #[test]
449 fn every_text_model_error_maps_one_to_one() {
450 use std::time::Duration;
451 assert_eq!(
452 model_error(TextModelError::NotConfigured),
453 ClassifierError::NotConfigured
454 );
455 assert_eq!(
456 model_error(TextModelError::Rejected("4xx words".to_owned())),
457 ClassifierError::Rejected("4xx words".to_owned())
458 );
459 assert_eq!(
460 model_error(TextModelError::Transient {
461 retry_after: Some(Duration::from_secs(7))
462 }),
463 ClassifierError::Transient {
464 retry_after: Some(Duration::from_secs(7))
465 }
466 );
467 assert_eq!(
468 model_error(TextModelError::Transient { retry_after: None }),
469 ClassifierError::Transient { retry_after: None }
470 );
471 assert_eq!(
472 model_error(TextModelError::Transport("boom".to_owned())),
473 ClassifierError::Transport("boom".to_owned())
474 );
475 assert_eq!(
478 model_error(TextModelError::Unsupported(
479 cratefield_core::Capability::Tools
480 )),
481 ClassifierError::NotConfigured
482 );
483 }
484
485 #[test]
486 fn the_schema_pins_each_offered_label_and_nothing_else() {
487 let questions = BTreeMap::from([
488 (
489 "topic".to_owned(),
490 Question::Choice {
491 instructions: "Which?".to_owned(),
492 criteria: BTreeMap::from([
493 ("billing".to_owned(), "money".to_owned()),
494 ("bugs".to_owned(), "broken".to_owned()),
495 ]),
496 },
497 ),
498 (
499 "angry".to_owned(),
500 Question::Noul {
501 instructions: "Angry?".to_owned(),
502 },
503 ),
504 ]);
505 let schema = answer_schema(&questions);
506 let answers = &schema["properties"]["answers"];
507 assert_eq!(answers["required"], serde_json::json!(["angry", "topic"]));
509 let topic = &answers["properties"]["topic"];
510 assert_eq!(
511 topic["properties"]["value"]["enum"],
512 serde_json::json!(["billing", "bugs"])
513 );
514 assert_eq!(
515 topic["properties"]["probabilities"]["required"],
516 serde_json::json!(["billing", "bugs"])
517 );
518 let angry = &answers["properties"]["angry"];
519 assert_eq!(
520 angry["properties"]["value"]["enum"],
521 serde_json::json!(["true", "false"])
522 );
523 let rendered = schema.to_string();
524 assert!(
525 !rendered.contains("spam"),
526 "an unoffered label is not pinned"
527 );
528 assert_eq!(schema["additionalProperties"], serde_json::json!(false));
529 }
530}