hyperhive/swarm-review-classifier/src/answer.rs

//! Reading the model's answer to one classification.
//!
//! An answer that does not parse, names an unknown verdict or carries a
//! certainty outside 0..=100 is an error, and the row it belongs to is
//! excluded rather than guessed at. A reason over the cap is not an error: it
//! is cut to the cap, so the verdict and certainty survive.

use serde::Deserialize;
use swarm_queue_client::review_classifier::Verdict;

/// The longest reason, in characters, an answer carries; a longer one is
/// truncated to this, ending in `…`.
pub const REASON_MAX_CHARS: usize = 500;

/// One validated answer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Answer {
    pub verdict: Verdict,
    pub certainty: u8,
    pub reason: String,
}

/// The answer's fields as the model wrote them, before validation.
#[derive(Deserialize)]
struct Raw {
    verdict: String,
    certainty: serde_json::Number,
    reason: String,
}

/// Parse `text`, the model's final message. A surrounding code fence is
/// tolerated; anything else around the object is not.
///
/// # Errors
/// A message describing what is wrong with the answer, without quoting it.
pub fn parse(text: &str) -> Result<Answer, String> {
    let body = unfence(text.trim());
    let raw: Raw = serde_json::from_str(body)
        .map_err(|e| format!("the answer is not the JSON asked for: {e}"))?;
    let verdict = match raw.verdict.as_str() {
        "match" => Verdict::Match,
        "no-match" => Verdict::NoMatch,
        "not-applicable" => Verdict::NotApplicable,
        other => {
            return Err(format!(
                "the answer's verdict {other:?} is not one of the three"
            ));
        }
    };
    let certainty = raw
        .certainty
        .as_u64()
        .and_then(|c| u8::try_from(c).ok())
        .filter(|c| *c <= 100)
        .ok_or_else(|| {
            format!(
                "the answer's certainty {} is not a whole 0-100",
                raw.certainty
            )
        })?;
    Ok(Answer {
        verdict,
        certainty,
        reason: cap_reason(raw.reason),
    })
}

/// `reason` cut to `REASON_MAX_CHARS` characters, the last being `…`, if longer.
fn cap_reason(reason: String) -> String {
    if reason.chars().count() <= REASON_MAX_CHARS {
        return reason;
    }
    let mut cut: String = reason.chars().take(REASON_MAX_CHARS - 1).collect();
    cut.push('…');
    cut
}

/// `text` without a surrounding code fence, with or without a language tag.
fn unfence(text: &str) -> &str {
    let Some(inner) = text.strip_prefix("```").and_then(|t| t.strip_suffix("```")) else {
        return text;
    };
    let inner = inner.strip_prefix("json").unwrap_or(inner);
    inner.trim()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_valid_answer_parses() {
        let a = parse(r#"{"verdict": "no-match", "certainty": 92, "reason": "Naive datetimes."}"#)
            .expect("valid");
        assert_eq!(
            a,
            Answer {
                verdict: Verdict::NoMatch,
                certainty: 92,
                reason: "Naive datetimes.".to_owned(),
            }
        );
        let fenced =
            "```json\n{\"verdict\":\"not-applicable\",\"certainty\":0,\"reason\":\"\"}\n```";
        assert_eq!(
            parse(fenced).expect("fenced").verdict,
            Verdict::NotApplicable
        );
    }

    #[test]
    fn a_malformed_answer_is_an_error() {
        for bad in [
            "",
            "yes",
            "Sure! {\"verdict\":\"match\",\"certainty\":90,\"reason\":\"x\"}",
            r#"{"verdict":"maybe","certainty":50,"reason":"x"}"#,
            r#"{"verdict":"match","certainty":101,"reason":"x"}"#,
            r#"{"verdict":"match","certainty":-1,"reason":"x"}"#,
            r#"{"verdict":"match","certainty":80.5,"reason":"x"}"#,
            r#"{"verdict":"match","reason":"x"}"#,
        ] {
            assert!(parse(bad).is_err(), "{bad:?} parsed");
        }
    }

    #[test]
    fn an_oversized_reason_is_truncated_and_one_at_the_cap_is_kept() {
        let answer = |n: usize| {
            format!(
                r#"{{"verdict":"match","certainty":90,"reason":"{}"}}"#,
                "é".repeat(n)
            )
        };
        // Characters, not bytes: `é` is two bytes.
        let at_cap = parse(&answer(REASON_MAX_CHARS)).expect("at the cap");
        assert_eq!(at_cap.reason, "é".repeat(REASON_MAX_CHARS));

        let long = parse(&answer(600)).expect("over the cap still parses");
        assert_eq!(long.verdict, Verdict::Match);
        assert_eq!(long.certainty, 90);
        assert_eq!(long.reason.chars().count(), REASON_MAX_CHARS);
        assert_eq!(
            long.reason,
            format!("{}…", "é".repeat(REASON_MAX_CHARS - 1))
        );
    }
}