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;

/// The length the prompt asks a reason to stay under. Below `REASON_MAX_CHARS`
/// so a model that overshoots the ask still lands inside the cap, uncut.
pub const REASON_ASKED_CHARS: usize = 350;
const _: () = assert!(REASON_ASKED_CHARS < REASON_MAX_CHARS);

/// 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 code fence around the object is
/// tolerated, with prose around the fence; prose around a bare 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 its code fence, with or without a language tag. A fence
/// around the whole message is stripped; otherwise the first fenced block in
/// the message is taken, so prose before or after it is dropped.
fn unfence(text: &str) -> &str {
    let inner = match text.strip_prefix("```").and_then(|t| t.strip_suffix("```")) {
        Some(inner) => inner,
        None => match text
            .split_once("```")
            .and_then(|(_, t)| t.split_once("```"))
        {
            Some((block, _)) => block,
            None => 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_fenced_answer_inside_prose_parses() {
        let a = parse(
            "Here is my answer:\n```json\n{\"verdict\":\"match\",\"certainty\":80,\"reason\":\"x\"}\n```\nHope that helps.",
        )
        .expect("fence in prose");
        assert_eq!((a.verdict, a.certainty), (Verdict::Match, 80));
        let untagged =
            "Answer: ```{\"verdict\":\"no-match\",\"certainty\":70,\"reason\":\"\"}``` done";
        assert_eq!(parse(untagged).expect("untagged").verdict, Verdict::NoMatch);
        assert!(parse("Here is my answer:\n```json\nnot json\n```").is_err());
        assert!(parse("an unclosed ```json\n{\"verdict\":\"match\"").is_err());
    }

    #[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))
        );
    }
}