gemini_adk_fluent_rs/compose/

judge.rs

//! LLM-as-judge — the shared async evaluation primitive behind the LLM-backed
//! `G::` guards and `E::` criteria.
//!
//! Mirrors ADK Python's `final_response_match_v2` / safety-evaluator approach: a
//! judge model is prompted to render a structured verdict, and the verdict label
//! is parsed back out of the model's reply (robust to surrounding prose).

use std::sync::Arc;

use gemini_adk_rs::llm::{BaseLlm, LlmRequest};
use gemini_genai_rs::prelude::{Content, Part, Role};

/// A configured LLM-as-judge. Cheap to clone (`Arc` inside).
#[derive(Clone)]
pub struct LlmJudge {
    llm: Arc<dyn BaseLlm>,
    /// Describes the condition that constitutes a *violation* (a `true` verdict).
    rubric: String,
    /// Label used for the supplied context/reference block in the prompt.
    context_label: &'static str,
    /// Whether to include the context/reference block at all.
    use_context: bool,
}

/// The judge's verdict.
pub struct Verdict {
    /// Whether the judge flagged a violation of the rubric.
    pub flagged: bool,
    /// Short reason / explanation extracted from the judge reply.
    pub reason: String,
}

impl LlmJudge {
    /// Create a judge for the given rubric (the condition that is a *violation*).
    pub fn new(llm: Arc<dyn BaseLlm>, rubric: impl Into<String>) -> Self {
        Self {
            llm,
            rubric: rubric.into(),
            context_label: "CONTEXT",
            use_context: false,
        }
    }

    /// Include a context/reference block in the judge prompt, under `label`
    /// (e.g. "CONTEXT" for grounding, "REFERENCE ANSWER" for semantic match).
    pub fn with_context(mut self, label: &'static str) -> Self {
        self.use_context = true;
        self.context_label = label;
        self
    }

    /// The rubric describing what counts as a violation.
    pub fn rubric(&self) -> &str {
        &self.rubric
    }

    fn prompt(&self, output: &str, context: Option<&str>) -> String {
        let mut p = String::with_capacity(output.len() + self.rubric.len() + 256);
        p.push_str(
            "You are a strict evaluator. Decide whether the RESPONSE below violates the \
             RULE. Judge only against the RULE; do not invent additional requirements.\n\n",
        );
        p.push_str("RULE — a violation is when the following is TRUE:\n");
        p.push_str(&self.rubric);
        p.push_str("\n\n");
        if self.use_context {
            if let Some(ctx) = context {
                p.push_str(self.context_label);
                p.push_str(":\n");
                p.push_str(ctx);
                p.push_str("\n\n");
            }
        }
        p.push_str("RESPONSE:\n");
        p.push_str(output);
        p.push_str(
            "\n\nReply with ONLY a single-line JSON object and nothing else:\n\
             {\"violation\": true|false, \"reason\": \"<at most 20 words>\"}",
        );
        p
    }

    /// Run the judge over an output (and optional context/reference).
    ///
    /// Fails open: if the judge LLM errors, the verdict is *not* flagged (so a
    /// transient judge outage never vetoes a turn) and the error is recorded in
    /// `reason`.
    pub async fn judge(&self, output: &str, context: Option<&str>) -> Verdict {
        let req = LlmRequest::from_contents(vec![Content::user(self.prompt(output, context))]);
        match self.llm.generate(req).await {
            Ok(resp) => parse_verdict(&resp.text()),
            Err(e) => Verdict {
                flagged: false,
                reason: format!("judge unavailable: {e}"),
            },
        }
    }
}

/// Parse a verdict from the judge model's reply. Tolerant of extra prose around
/// the JSON: it scans for the `violation` field's boolean and the `reason`
/// string, falling back to common labels (`invalid`, `unsafe`).
pub fn parse_verdict(text: &str) -> Verdict {
    let lower = text.to_ascii_lowercase();
    let flagged = match lower.find("violation") {
        Some(idx) => {
            let tail = &lower[idx..];
            match (tail.find("true"), tail.find("false")) {
                (Some(t), Some(f)) => t < f,
                (Some(_), None) => true,
                (None, Some(_)) => false,
                (None, None) => lower.contains("invalid") || lower.contains("unsafe"),
            }
        }
        // No explicit `violation` field — fall back to common labels. Note
        // `contains("invalid")` is false for "valid" but true for "INVALID".
        None => lower.contains("invalid") || lower.contains("unsafe"),
    };
    let reason = extract_reason(text).unwrap_or_else(|| {
        if flagged {
            "flagged by judge".to_string()
        } else {
            "ok".to_string()
        }
    });
    Verdict { flagged, reason }
}

fn extract_reason(text: &str) -> Option<String> {
    let key = "\"reason\"";
    let i = text.find(key)?;
    let after = &text[i + key.len()..];
    let colon = after.find(':')?;
    let rest = after[colon + 1..].trim_start();
    let rest = rest.strip_prefix('"')?;
    let end = rest.find('"')?;
    Some(rest[..end].to_string())
}

/// Render conversation history into a plain-text block for a judge prompt,
/// keeping role labels so the judge can reason about grounding.
pub fn render_contents(contents: &[Content]) -> String {
    let mut out = String::new();
    for content in contents {
        let role = match content.role {
            Some(Role::User) => "user",
            Some(Role::Model) => "model",
            _ => "system",
        };
        for part in &content.parts {
            if let Part::Text { text } = part {
                if !text.is_empty() {
                    out.push_str(role);
                    out.push_str(": ");
                    out.push_str(text);
                    out.push('\n');
                }
            }
        }
    }
    out
}

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

    #[test]
    fn parses_violation_true() {
        let v = parse_verdict(r#"{"violation": true, "reason": "contains slur"}"#);
        assert!(v.flagged);
        assert_eq!(v.reason, "contains slur");
    }

    #[test]
    fn parses_violation_false() {
        let v = parse_verdict("Sure! {\"violation\": false, \"reason\": \"all good\"}");
        assert!(!v.flagged);
        assert_eq!(v.reason, "all good");
    }

    #[test]
    fn falls_back_to_labels() {
        assert!(parse_verdict("Verdict: INVALID").flagged);
        assert!(!parse_verdict("looks valid to me").flagged);
    }
}