clawmates/crates/cm-orchestrator/src/judge.rs

//! An LLM-judge [`Scorer`]: ask a model to rate how well a run's output
//! accomplishes the task, on a 0–100 scale, normalized to `[0,1]`.
//!
//! Gives the comparison harness real quality numbers (vs the deterministic
//! scorers used in tests). Behind the `provider` feature.

use std::sync::Arc;

use cm_llm::{ChatMessage, ChatRequest, ChatRole, ContentPart, LlmEvent, LlmProvider};
use futures::StreamExt;

use crate::{RunRecord, Scorer};

/// Scores run quality with an LLM via any `cm-llm` provider.
pub struct JudgeScorer {
    provider: Arc<dyn LlmProvider>,
    model: String,
    max_tokens: u32,
}

impl JudgeScorer {
    /// Build a judge over a shared provider, model id, and token budget.
    pub fn new(provider: Arc<dyn LlmProvider>, model: impl Into<String>, max_tokens: u32) -> Self {
        JudgeScorer {
            provider,
            model: model.into(),
            max_tokens,
        }
    }
}

/// Extract the first integer in the text and map 0..100 → 0.0..1.0.
/// Returns 0.0 if no integer is present.
fn parse_score(s: &str) -> f64 {
    let mut digits = String::new();
    for c in s.chars() {
        if c.is_ascii_digit() {
            digits.push(c);
        } else if !digits.is_empty() {
            break;
        }
    }
    digits
        .parse::<f64>()
        .ok()
        .map(|n| (n / 100.0).clamp(0.0, 1.0))
        .unwrap_or(0.0)
}

impl Scorer for JudgeScorer {
    async fn score(&self, task: &str, record: &RunRecord) -> f64 {
        let request = ChatRequest {
            system: "You are a strict evaluator. Rate how well RESULT accomplishes TASK on a \
                     scale from 0 to 100. Reply with ONLY the integer."
                .to_string(),
            messages: vec![ChatMessage {
                role: ChatRole::User,
                parts: vec![ContentPart::text(format!(
                    "TASK:\n{task}\n\nRESULT:\n{}",
                    record.final_output
                ))],
            }],
            tools: vec![],
            model: self.model.clone(),
            max_tokens: self.max_tokens,
            web_search: false,
        };

        let mut stream = match self.provider.stream(request).await {
            Ok(s) => s,
            Err(_) => return 0.0,
        };
        let mut text = String::new();
        while let Some(event) = stream.next().await {
            if let Ok(LlmEvent::TextDelta(t)) = event {
                text.push_str(&t);
            }
        }
        parse_score(&text)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{RunMetrics, RunRecord};
    use cm_topology::TopologyKind;

    fn record(output: &str) -> RunRecord {
        RunRecord {
            kind: TopologyKind::Flat,
            steps: vec![],
            final_output: output.to_string(),
            totals: RunMetrics::default(),
        }
    }

    #[test]
    fn parse_score_handles_text_and_garbage() {
        assert!((parse_score("85") - 0.85).abs() < 1e-9);
        assert!((parse_score("Score: 92/100") - 0.92).abs() < 1e-9);
        assert!((parse_score("over") - 0.0).abs() < 1e-9);
        assert!((parse_score("150") - 1.0).abs() < 1e-9); // clamped
    }

    #[tokio::test]
    async fn judge_returns_normalized_score() {
        // Scenario keyed on "RESULT" (always in the judge prompt) returns "85".
        let toml = r#"
            [[scenario]]
            marker = "RESULT"
            [[scenario.turns]]
            [[scenario.turns.events]]
            type = "text"
            text = "85"
        "#;
        let provider = Arc::new(cm_llm::ScriptedProvider::from_toml(toml).unwrap());
        let judge = JudgeScorer::new(provider, "test-model", 16);
        let score = judge.score("write a poem", &record("a fine poem")).await;
        assert!((score - 0.85).abs() < 1e-9, "got {score}");
    }
}