hyperhive/swarm-queue-client/src/review_classifier.rs

//! The review classifier's runs: the subject swarm-controller sends a run
//! on, and what travels each way. A try-out and a requested review are the
//! same request.
//!
//! Three crates name these: the controller sends the request, the classifier
//! answers it, and the auth-callout responder grants the one subject to each.

use std::time::Duration;

use serde::{Deserialize, Serialize};

use crate::agent_inference::Effort;

/// The subject every run goes to. The controller is its one
/// publisher and the classifier its one subscriber; the classifier answers
/// on the request's reply subject.
pub const SUBJECT: &str = "$SWARM.review-classifier.try-out";

/// The queue group the classifier subscribes in, so a second instance shares
/// the requests instead of running each one twice.
pub const QUEUE_GROUP: &str = "review-classifier";

/// The most classifications one run carries.
pub const MAX_CLASSIFICATIONS: usize = 10;

/// The longest one classification's model run may take before it is stopped
/// and its row carries a timeout.
pub const CLASSIFICATION_DEADLINE: Duration = Duration::from_mins(3);

/// The longest a whole run may take: every classification at its deadline,
/// run one after another, plus a minute for the rest. The controller waits
/// this long for the reply, and the queue lets the classifier answer for
/// this long.
pub const RUN_DEADLINE: Duration = Duration::from_mins(31);

/// The model a run uses when its request names none, and the swarm-wide
/// setting's default.
pub const DEFAULT_MODEL: &str = "haiku";

fn default_model() -> String {
    DEFAULT_MODEL.to_owned()
}

/// One run: the PR as the controller read it, and the classifications to
/// answer about it.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TryOut {
    /// Who started the run: the swarm UI user, as the gateway named them, or
    /// the forge user who requested the classifier's review.
    pub operator: String,
    /// `owner/name`.
    pub repo: String,
    pub pr: u64,
    /// The commit the commits and the diff were read at.
    pub head_sha: String,
    /// Oldest first.
    pub commits: Vec<Commit>,
    /// The PR's unified diff, as the forge serves it.
    pub diff: String,
    pub classifications: Vec<Classification>,
    /// The model every classification runs on, as claude's `--model` takes it.
    /// Absent is [`DEFAULT_MODEL`], so a controller that sends none still gets
    /// its run answered.
    #[serde(default = "default_model")]
    pub model: String,
    /// The effort every classification runs at; `None` passes no `--effort`.
    pub effort: Option<Effort>,
    /// What every classification runs on. Absent is [`Runtime::ApiKey`], so a
    /// controller that sends none still gets its run answered.
    #[serde(default)]
    pub runtime: Runtime,
}

/// What a run's classifications run on.
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "kebab-case")]
pub enum Runtime {
    /// `claude` on the classifier's Anthropic API key, at [`TryOut::model`]
    /// and [`TryOut::effort`].
    #[default]
    ApiKey,
    /// An ACP agent on the classifier's ACP provider key. [`TryOut::model`]
    /// is claude's and unused here; [`TryOut::effort`] applies where the
    /// agent offers it.
    Acp(Acp),
}

/// The ACP provider a run's classifications go to.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct Acp {
    /// The provider's OpenAI-compatible API root, `https://` only.
    pub base_url: String,
    /// The model, as the provider names it.
    pub model: String,
    pub prices: Prices,
}

/// What the provider charges, in US dollars per million tokens. A run's cost
/// is computed from these and the tokens the agent reports.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
#[serde(deny_unknown_fields)]
pub struct Prices {
    pub input: f64,
    pub output: f64,
    /// Absent prices cache reads as input.
    #[serde(default)]
    pub cache_read: Option<f64>,
    /// Absent prices cache writes as input.
    #[serde(default)]
    pub cache_write: Option<f64>,
}

/// One commit of the PR.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Commit {
    pub sha: String,
    /// The full message, subject and body.
    pub message: String,
}

/// One yes/no rule to check the PR against.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
pub struct Classification {
    pub name: String,
    /// The yes/no question the model answers about the PR.
    pub description: String,
    pub polarity: Polarity,
    /// The certainty, 0 to 100, a verdict needs to count.
    pub threshold: u8,
}

/// Which verdict is the wanted one.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
#[serde(rename_all = "lowercase")]
pub enum Polarity {
    /// A match is the wanted outcome.
    Positive,
    /// A match is a violation.
    Negative,
}

/// The classifier's answer to one [`TryOut`].
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "state", rename_all = "kebab-case")]
pub enum Reply {
    Done {
        rows: Vec<Row>,
        overall: Overall,
    },
    /// Nothing ran to a verdict: no API key, or every classification failed.
    Failed {
        error: String,
    },
}

/// One classification's result, as the result table shows it.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
pub struct Row {
    pub name: String,
    pub polarity: Polarity,
    pub verdict: Verdict,
    pub certainty: u8,
    pub threshold: u8,
    /// Whether the row counts toward the overall verdict: applicable, and
    /// its certainty reached the threshold.
    pub counted: bool,
    /// For a counted row, whether the verdict is on the wanted side of the
    /// polarity. `false` for every row that is not counted.
    pub passed: bool,
    pub reason: String,
    /// Why this classification has no answer: a runtime error, a timeout, or
    /// an answer that did not validate. Such a row reads `not-applicable`,
    /// certainty 0 and an empty reason, and is never counted; key off this
    /// field, not off the verdict.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
#[serde(rename_all = "kebab-case")]
pub enum Verdict {
    Match,
    NoMatch,
    NotApplicable,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "schema", derive(utoipa::ToSchema))]
#[serde(rename_all = "kebab-case")]
pub enum Overall {
    Approve,
    RequestChanges,
}

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

    #[test]
    fn the_wire_spells_the_enums_as_the_api_does() {
        let row = Row {
            name: "datetime".to_owned(),
            polarity: Polarity::Positive,
            verdict: Verdict::NotApplicable,
            certainty: 95,
            threshold: 70,
            counted: false,
            passed: false,
            reason: "The PR adds no timestamp.".to_owned(),
            error: None,
        };
        let json = serde_json::to_value(Reply::Done {
            rows: vec![row],
            overall: Overall::RequestChanges,
        })
        .expect("encodes");
        assert_eq!(json["state"], "done");
        assert_eq!(json["overall"], "request-changes");
        assert_eq!(json["rows"][0]["polarity"], "positive");
        assert_eq!(json["rows"][0]["verdict"], "not-applicable");
        assert!(json["rows"][0].get("error").is_none());
    }

    #[test]
    fn the_run_deadline_covers_every_classification_at_its_deadline() {
        let all = CLASSIFICATION_DEADLINE * u32::try_from(MAX_CLASSIFICATIONS).expect("small");
        assert!(RUN_DEADLINE > all);
    }

    #[test]
    fn a_request_without_a_model_runs_on_the_default() {
        let mut request = serde_json::json!({
            "operator": "alice",
            "repo": "o/r",
            "pr": 7,
            "head_sha": "abc123",
            "commits": [],
            "diff": "",
            "classifications": [],
        });
        let run: TryOut = serde_json::from_value(request.clone()).expect("decodes");
        assert_eq!((run.model.as_str(), run.effort), (DEFAULT_MODEL, None));
        assert_eq!(run.runtime, Runtime::ApiKey);
        // Control: a named model and effort are the ones used.
        request["model"] = "sonnet".into();
        request["effort"] = "high".into();
        let run: TryOut = serde_json::from_value(request).expect("decodes");
        assert_eq!(
            (run.model.as_str(), run.effort),
            ("sonnet", Some(Effort::High))
        );
    }

    #[test]
    fn an_acp_runtime_is_tagged_by_kind() {
        let acp = Runtime::Acp(Acp {
            base_url: "https://llm.example/v1".to_owned(),
            model: "qwen".to_owned(),
            prices: Prices {
                input: 0.5,
                output: 2.0,
                cache_read: None,
                cache_write: None,
            },
        });
        let json = serde_json::to_value(&acp).expect("encodes");
        assert_eq!(json["kind"], "acp");
        assert_eq!(json["prices"]["input"], 0.5);
        assert_eq!(
            serde_json::from_value::<Runtime>(json).expect("decodes"),
            acp
        );
        assert_eq!(
            serde_json::to_value(Runtime::ApiKey).expect("encodes"),
            serde_json::json!({"kind": "api-key"})
        );
    }
}