use std::collections::BTreeMap;
use std::time::Duration;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
pub const DEFAULT_MAX_STATE_CHARS: usize = 96_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Calibration {
Classifier,
LanguageModel,
}
impl Calibration {
pub fn name(&self) -> &'static str {
match self {
Calibration::Classifier => "classifier",
Calibration::LanguageModel => "language_model",
}
}
}
impl std::fmt::Display for Calibration {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.name())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct ClassifierProfile {
pub calibration: Calibration,
pub max_state_chars: usize,
}
impl ClassifierProfile {
#[must_use]
pub fn new(calibration: Calibration, max_state_chars: usize) -> Self {
Self {
calibration,
max_state_chars,
}
}
#[must_use]
pub fn truncate<'a>(&self, state: &'a str) -> (&'a str, bool) {
if state.len() <= self.max_state_chars {
return (state, false);
}
let mut end = self.max_state_chars;
while end > 0 && !state.is_char_boundary(end) {
end -= 1;
}
(&state[..end], true)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Question {
Choice {
instructions: String,
criteria: BTreeMap<String, String>,
},
Score {
instructions: String,
levels: Vec<(String, String)>,
},
Noul { instructions: String },
}
impl Question {
#[must_use]
pub fn instructions(&self) -> &str {
match self {
Question::Choice { instructions, .. }
| Question::Score { instructions, .. }
| Question::Noul { instructions } => instructions,
}
}
#[must_use]
pub fn kind(&self) -> &'static str {
match self {
Question::Choice { .. } => "choice",
Question::Score { .. } => "score",
Question::Noul { .. } => "noul",
}
}
#[must_use]
pub fn labels(&self) -> Vec<&str> {
match self {
Question::Choice { criteria, .. } => criteria.keys().map(String::as_str).collect(),
Question::Score { levels, .. } => {
levels.iter().map(|(name, _)| name.as_str()).collect()
}
Question::Noul { .. } => vec!["true", "false"],
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum AnswerValue {
Choice(String),
Score(f32),
Noul(bool),
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct Answer {
pub value: AnswerValue,
pub probabilities: BTreeMap<String, f32>,
pub confidence: f32,
pub model: Option<String>,
}
impl Answer {
#[must_use]
pub fn new(value: AnswerValue, probabilities: BTreeMap<String, f32>, confidence: f32) -> Self {
Self {
value,
probabilities,
confidence,
model: None,
}
}
#[must_use]
pub fn with_model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
#[must_use]
pub fn choice(label: impl Into<String>, probabilities: BTreeMap<String, f32>) -> Self {
let label = label.into();
let confidence = probabilities.get(&label).copied().unwrap_or(0.0);
Self {
value: AnswerValue::Choice(label),
probabilities,
confidence,
model: None,
}
}
#[must_use]
pub fn score(value: f32, probabilities: BTreeMap<String, f32>) -> Self {
let confidence = probabilities
.get(&value.to_string())
.copied()
.unwrap_or(0.0);
Self {
value: AnswerValue::Score(value),
probabilities,
confidence,
model: None,
}
}
#[must_use]
pub fn noul(value: bool, probabilities: BTreeMap<String, f32>) -> Self {
let label = if value { "true" } else { "false" };
let confidence = probabilities.get(label).copied().unwrap_or(0.0);
Self {
value: AnswerValue::Noul(value),
probabilities,
confidence,
model: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClassifierError {
NotConfigured,
Rejected(String),
Transient { retry_after: Option<Duration> },
Transport(String),
}
impl std::fmt::Display for ClassifierError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let scrub = crate::logging::scrub_text;
match self {
Self::NotConfigured => f.write_str("no classifier is wired for this venture"),
Self::Rejected(message) => write!(f, "classification rejected: {}", scrub(message)),
Self::Transient { .. } => f.write_str("classification failed, retryable"),
Self::Transport(message) => {
write!(f, "classification transport failed: {}", scrub(message))
}
}
}
}
impl std::error::Error for ClassifierError {}
impl ClassifierError {
pub fn retry_after(&self) -> Option<Duration> {
match self {
ClassifierError::Transient { retry_after } => *retry_after,
_ => None,
}
}
}
pub fn validate_questions(questions: &BTreeMap<String, Question>) -> Result<(), ClassifierError> {
if questions.is_empty() {
return Err(ClassifierError::Rejected(
"a question set must not be empty".to_owned(),
));
}
for (id, question) in questions {
if id.trim().is_empty() {
return Err(ClassifierError::Rejected(format!(
"question id must not be blank, got {id:?}"
)));
}
if question.instructions().trim().is_empty() {
return Err(ClassifierError::Rejected(format!(
"question {id:?} has blank instructions"
)));
}
match question {
Question::Choice { criteria, .. } if criteria.len() < 2 => {
return Err(ClassifierError::Rejected(format!(
"choice question {id:?} needs at least two criteria, got {}",
criteria.len()
)));
}
Question::Score { levels, .. } if levels.len() < 2 => {
return Err(ClassifierError::Rejected(format!(
"score question {id:?} needs at least two levels, got {}",
levels.len()
)));
}
_ => {}
}
}
Ok(())
}
#[async_trait]
pub trait Classifier: Send + Sync {
fn profile(&self) -> ClassifierProfile;
async fn ask(
&self,
state: &str,
questions: &BTreeMap<String, Question>,
) -> Result<BTreeMap<String, Answer>, ClassifierError>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_calibration_names_itself_and_round_trips_through_json() {
for calibration in [Calibration::Classifier, Calibration::LanguageModel] {
let json = serde_json::to_string(&calibration).expect("serialises");
let back: Calibration = serde_json::from_str(&json).expect("deserialises");
assert_eq!(calibration, back);
}
assert_eq!(
serde_json::to_value(Calibration::LanguageModel).unwrap(),
"language_model"
);
assert_eq!(Calibration::Classifier.name(), "classifier");
assert_eq!(Calibration::LanguageModel.to_string(), "language_model");
}
#[test]
fn truncate_is_a_no_op_under_the_limit() {
assert_eq!(DEFAULT_MAX_STATE_CHARS, 96_000);
let profile = ClassifierProfile::new(Calibration::Classifier, DEFAULT_MAX_STATE_CHARS);
let state = "a state well under the ceiling";
let (sent, dropped) = profile.truncate(state);
assert_eq!(sent, state);
assert!(!dropped);
}
#[test]
fn truncate_trims_deterministically_over_the_limit() {
let profile = ClassifierProfile::new(Calibration::Classifier, 16);
let state = "the same state, asked twice, over the limit";
let (first, dropped) = profile.truncate(state);
let (second, dropped_again) = profile.truncate(state);
assert!(dropped);
assert!(dropped_again);
assert_eq!(first, second, "the same state always trims the same way");
assert_eq!(first.len(), 16, "the prefix spends the whole budget");
assert!(state.starts_with(first), "the prefix is a prefix");
}
#[test]
fn truncate_never_splits_a_multi_byte_char() {
let state = "漢字漢字";
let profile = ClassifierProfile::new(Calibration::LanguageModel, 7);
let (sent, dropped) = profile.truncate(state);
assert!(dropped);
assert_eq!(sent, "漢字");
assert_eq!(sent.chars().count(), 2);
}
#[test]
fn a_zero_limit_keeps_nothing_but_stays_on_a_boundary() {
let profile = ClassifierProfile::new(Calibration::Classifier, 0);
let (sent, dropped) = profile.truncate("漢字");
assert_eq!(sent, "");
assert!(dropped);
}
#[test]
fn a_question_names_its_shape_and_labels() {
let choice = Question::Choice {
instructions: "Pick.".to_owned(),
criteria: BTreeMap::from([
("billing".to_owned(), "money and invoices".to_owned()),
("bugs".to_owned(), "something is broken".to_owned()),
]),
};
assert_eq!(choice.kind(), "choice");
assert_eq!(choice.instructions(), "Pick.");
assert_eq!(choice.labels(), vec!["billing", "bugs"]);
let score = Question::Score {
instructions: "Rate.".to_owned(),
levels: vec![
("1".to_owned(), "low".to_owned()),
("5".to_owned(), "high".to_owned()),
],
};
assert_eq!(score.kind(), "score");
assert_eq!(score.labels(), vec!["1", "5"]);
let noul = Question::Noul {
instructions: "Yes or no?".to_owned(),
};
assert_eq!(noul.kind(), "noul");
assert_eq!(noul.instructions(), "Yes or no?");
assert_eq!(noul.labels(), vec!["true", "false"]);
}
#[test]
fn answer_conveniences_pick_confidence_from_the_probabilities() {
let criteria = BTreeMap::from([("spam".to_owned(), 0.9_f32), ("ham".to_owned(), 0.1_f32)]);
let answer = Answer::choice("spam", criteria.clone());
assert_eq!(answer.value, AnswerValue::Choice("spam".to_owned()));
assert_eq!(answer.probabilities, criteria);
assert!((answer.confidence - 0.9).abs() < f32::EPSILON);
let levels = BTreeMap::from([("1".to_owned(), 0.25_f32), ("2".to_owned(), 0.75_f32)]);
let answer = Answer::score(2.0, levels);
assert_eq!(answer.value, AnswerValue::Score(2.0));
assert!((answer.confidence - 0.75).abs() < f32::EPSILON);
let verdicts =
BTreeMap::from([("true".to_owned(), 0.7_f32), ("false".to_owned(), 0.3_f32)]);
let answer = Answer::noul(true, verdicts);
assert_eq!(answer.value, AnswerValue::Noul(true));
assert!((answer.confidence - 0.7).abs() < f32::EPSILON);
}
#[test]
fn an_absent_label_yields_zero_confidence() {
let verdicts = BTreeMap::from([("true".to_owned(), 1.0_f32)]);
assert!(
Answer::noul(false, verdicts).confidence.abs() < f32::EPSILON,
"an absent label answers 0.0"
);
let confidence = Answer::choice(
"unheard-of",
BTreeMap::from([("known".to_owned(), 1.0_f32)]),
)
.confidence;
assert!(
confidence.abs() < f32::EPSILON,
"an absent label answers 0.0"
);
assert!(
Answer::score(3.5, BTreeMap::new()).confidence.abs() < f32::EPSILON,
"an absent label answers 0.0"
);
let answer = Answer::new(AnswerValue::Score(4.0), BTreeMap::new(), 0.5);
assert!((answer.confidence - 0.5).abs() < f32::EPSILON);
}
#[test]
fn with_model_records_the_reporting_version_and_defaults_to_none() {
let answer = Answer::noul(true, BTreeMap::new());
assert_eq!(answer.model, None);
assert_eq!(
Answer::choice("spam", BTreeMap::new())
.with_model("jev-1.13.0")
.model
.as_deref(),
Some("jev-1.13.0")
);
assert_eq!(
Answer::score(3.0, BTreeMap::new()).with_model("v2").model,
Some("v2".to_owned())
);
}
fn good_set() -> BTreeMap<String, Question> {
BTreeMap::from([
(
"topic".to_owned(),
Question::Choice {
instructions: "Which topic does the note belong to?".to_owned(),
criteria: BTreeMap::from([
("billing".to_owned(), "money, invoices".to_owned()),
("bugs".to_owned(), "something is broken".to_owned()),
]),
},
),
(
"severity".to_owned(),
Question::Score {
instructions: "How severe is the note?".to_owned(),
levels: vec![
("1".to_owned(), "a typo".to_owned()),
("5".to_owned(), "data loss".to_owned()),
],
},
),
(
"angry".to_owned(),
Question::Noul {
instructions: "Is the writer angry?".to_owned(),
},
),
])
}
#[test]
fn a_good_set_is_accepted() {
assert_eq!(validate_questions(&good_set()), Ok(()));
}
#[test]
fn an_empty_set_is_rejected() {
let error = validate_questions(&BTreeMap::new()).unwrap_err();
assert!(matches!(error, ClassifierError::Rejected(_)));
}
#[test]
fn a_blank_id_is_rejected() {
let mut questions = good_set();
questions.insert(" ".to_owned(), questions["topic"].clone());
let error = validate_questions(&questions).unwrap_err();
assert!(matches!(error, ClassifierError::Rejected(_)));
}
#[test]
fn blank_instructions_are_rejected() {
let questions = BTreeMap::from([(
"topic".to_owned(),
Question::Noul {
instructions: " ".to_owned(),
},
)]);
assert!(matches!(
validate_questions(&questions),
Err(ClassifierError::Rejected(_))
));
}
#[test]
fn a_choice_with_one_criterion_is_rejected() {
let mut questions = good_set();
questions.insert(
"solo".to_owned(),
Question::Choice {
instructions: "Pick one.".to_owned(),
criteria: BTreeMap::from([("only".to_owned(), "the only one".to_owned())]),
},
);
assert!(matches!(
validate_questions(&questions),
Err(ClassifierError::Rejected(_))
));
}
#[test]
fn a_score_with_one_level_is_rejected() {
let mut questions = good_set();
questions.insert(
"one-note".to_owned(),
Question::Score {
instructions: "Rate it.".to_owned(),
levels: vec![("3".to_owned(), "middling".to_owned())],
},
);
assert!(matches!(
validate_questions(&questions),
Err(ClassifierError::Rejected(_))
));
}
#[test]
fn display_sanitizes_the_provider_text() {
let error = ClassifierError::Rejected(
"provider 400 for https://api.example.test/v1/classify?token=secret-abcdef".to_owned(),
);
let text = error.to_string();
assert!(text.contains("classification rejected"), "{text}");
assert!(!text.contains("secret-abcdef"), "{text}");
assert!(text.contains("?[redacted]"), "{text}");
let error = ClassifierError::Transport("timeout quoting alice@example.test".to_owned());
let text = error.to_string();
assert!(text.contains("classification transport failed"), "{text}");
assert!(!text.contains('@'), "{text}");
assert!(format!("{error:?}").contains("alice@example.test"));
}
#[test]
fn transient_says_retryable_and_carries_no_text_to_scrub() {
let error = ClassifierError::Transient {
retry_after: Some(Duration::from_secs(30)),
};
assert_eq!(error.to_string(), "classification failed, retryable");
}
#[test]
fn not_configured_names_the_missing_wiring_and_carries_no_back_off() {
assert_eq!(
ClassifierError::NotConfigured.to_string(),
"no classifier is wired for this venture"
);
assert_eq!(ClassifierError::NotConfigured.retry_after(), None);
assert_eq!(
ClassifierError::Rejected("422".to_owned()).retry_after(),
None,
"a rejection is not a back-off"
);
assert_eq!(
ClassifierError::Transport("timed out".to_owned()).retry_after(),
None
);
let throttled = ClassifierError::Transient {
retry_after: Some(Duration::from_secs(30)),
};
assert_eq!(throttled.retry_after(), Some(Duration::from_secs(30)));
}
}