aether_evals/
judge.rs

1use crate::evals::format_transcript;
2use aether_core::events::AgentMessage;
3use futures::StreamExt;
4use llm::types::IsoString;
5use llm::{ChatMessage, ContentBlock, Context, LlmResponse, StreamingModelProvider};
6use schemars::{JsonSchema, Schema, schema_for};
7use serde::{Deserialize, Serialize};
8use std::borrow::Borrow;
9use std::collections::{BTreeMap, BTreeSet};
10use thiserror::Error;
11
12/// Start building an LLM-as-judge from structured context and rubric criteria.
13pub fn judge() -> JudgeBuilder {
14    JudgeBuilder::default()
15}
16
17/// A built judge: the assembled prompt plus the normalized rubric it grades against. Run it with
18/// [`Judge::run`] against a model, or grade a parsed [`JudgeRubricResponse`] with
19/// [`Judge::summarize`].
20#[derive(Debug, Clone)]
21pub struct Judge {
22    pub prompt: String,
23    pub criteria: Vec<JudgeCriterionSpec>,
24}
25
26#[derive(Debug, Clone, Default)]
27pub struct JudgeBuilder {
28    instructions: Option<String>,
29    task: Option<String>,
30    context: JudgeContext,
31    criteria: Vec<JudgeCriterionSpec>,
32}
33
34/// Evidence the judge grades against: the agent transcript, a workspace diff, and/or final files.
35#[derive(Debug, Clone, Default)]
36pub struct JudgeContext {
37    pub transcript: Option<Vec<AgentMessage>>,
38    pub diff: Option<String>,
39    pub files: BTreeMap<String, String>,
40}
41
42/// A single rubric criterion scored on a normalized 0.0..=1.0 scale.
43#[derive(Debug, Clone, JsonSchema)]
44#[serde(rename_all = "camelCase", deny_unknown_fields)]
45pub struct JudgeCriterionSpec {
46    pub id: String,
47    pub description: String,
48    #[serde(default = "default_blocking")]
49    pub blocking: bool,
50    #[serde(default = "default_weight")]
51    pub weight: f64,
52    #[serde(default = "default_threshold")]
53    pub threshold: f64,
54}
55
56/// The graded result of running a judge: an overall pass/score plus per-criterion detail.
57#[derive(Debug, Clone, Serialize, JsonSchema)]
58#[serde(rename_all = "camelCase", deny_unknown_fields)]
59pub struct JudgeSummary {
60    pub passed: bool,
61    pub score: f64,
62    pub reason: String,
63    pub criteria: Vec<JudgeCriterionSummary>,
64}
65
66#[derive(Debug, Clone, Serialize, JsonSchema)]
67#[serde(rename_all = "camelCase", deny_unknown_fields)]
68pub struct JudgeCriterionSummary {
69    pub id: String,
70    pub description: String,
71    pub blocking: bool,
72    pub weight: f64,
73    pub threshold: f64,
74    pub score: f64,
75    pub passed: bool,
76    pub reason: String,
77}
78
79/// The raw rubric response the judge model is expected to return.
80#[derive(Debug, Deserialize, JsonSchema)]
81#[serde(deny_unknown_fields)]
82pub struct JudgeRubricResponse {
83    pub criteria: Vec<JudgeCriterionResponse>,
84    pub overall_reason: String,
85}
86
87#[derive(Debug, Deserialize, JsonSchema)]
88#[serde(deny_unknown_fields)]
89pub struct JudgeCriterionResponse {
90    pub id: String,
91    pub score: f64,
92    pub reason: String,
93}
94
95#[derive(Debug, Error)]
96pub enum JudgeError {
97    #[error("invalid judge input: {0}")]
98    InvalidInput(String),
99
100    #[error("judge LLM stream error: {0}")]
101    Stream(#[from] llm::LlmError),
102
103    #[error("judge returned invalid JSON: {source}\nRaw response: {raw_response}")]
104    InvalidJson {
105        #[source]
106        source: serde_json::Error,
107        raw_response: String,
108    },
109
110    #[error("judge returned invalid judgment: {reason}\nRaw response: {raw_response}")]
111    InvalidJudgment { reason: String, raw_response: String },
112}
113
114impl Judge {
115    pub fn response_schema() -> Schema {
116        JudgeRubricResponse::schema()
117    }
118
119    /// Grade `llm` against this judge's rubric: stream the model's response, parse it as a
120    /// [`JudgeRubricResponse`], and summarize it.
121    pub async fn run(&self, llm: &dyn StreamingModelProvider) -> Result<JudgeSummary, JudgeError> {
122        tracing::info!("Running LLM judge");
123        let raw_response = self.stream_response(llm).await?;
124        let response: JudgeRubricResponse = serde_json::from_str(extract_json_object(&raw_response))
125            .map_err(|source| JudgeError::InvalidJson { source, raw_response: raw_response.clone() })?;
126        self.summarize(response)
127    }
128
129    pub fn summarize(&self, response: JudgeRubricResponse) -> Result<JudgeSummary, JudgeError> {
130        let mut responses = BTreeMap::new();
131        for criterion in response.criteria {
132            let id = criterion.id.clone();
133            if responses.insert(id.clone(), criterion).is_some() {
134                return Err(invalid_judgment(format!("duplicate response criterion id `{id}`"), ""));
135            }
136        }
137
138        let mut summaries = Vec::with_capacity(self.criteria.len());
139        let mut weighted_score = 0.0;
140        let mut total_weight = 0.0;
141        let mut blocking_failed = false;
142
143        for criterion in &self.criteria {
144            let Some(response) = responses.remove(&criterion.id) else {
145                return Err(invalid_judgment(format!("missing response criterion `{}`", criterion.id), ""));
146            };
147            if !response.score.is_finite() || !(0.0..=1.0).contains(&response.score) {
148                return Err(invalid_judgment(
149                    format!("criterion `{}` score must be between 0.0 and 1.0", criterion.id),
150                    "",
151                ));
152            }
153
154            let passed = response.score >= criterion.threshold;
155            blocking_failed |= criterion.blocking && !passed;
156            weighted_score += response.score * criterion.weight;
157            total_weight += criterion.weight;
158            summaries.push(JudgeCriterionSummary {
159                id: criterion.id.clone(),
160                description: criterion.description.clone(),
161                blocking: criterion.blocking,
162                weight: criterion.weight,
163                threshold: criterion.threshold,
164                score: response.score,
165                passed,
166                reason: response.reason,
167            });
168        }
169
170        if let Some(id) = responses.keys().next() {
171            return Err(invalid_judgment(format!("unknown response criterion `{id}`"), ""));
172        }
173
174        let weighted_score = weighted_score / total_weight;
175        let score = if blocking_failed { 0.0 } else { weighted_score };
176        let reason = if blocking_failed {
177            format!("weighted score {:.2}; one or more blockers failed; {}", weighted_score, response.overall_reason)
178        } else {
179            format!("weighted score {:.2}; all blockers met; {}", weighted_score, response.overall_reason)
180        };
181
182        Ok(JudgeSummary { passed: !blocking_failed, score, reason, criteria: summaries })
183    }
184
185    async fn stream_response(&self, llm: &dyn StreamingModelProvider) -> Result<String, JudgeError> {
186        let message =
187            ChatMessage::User { content: vec![ContentBlock::text(self.prompt.clone())], timestamp: IsoString::now() };
188        let mut response_stream = llm.stream_response(&Context::new(vec![message], vec![]));
189        let mut raw_response = String::new();
190        while let Some(result) = response_stream.next().await {
191            match result {
192                Ok(LlmResponse::Text { chunk }) => raw_response.push_str(&chunk),
193                Err(error) => return Err(JudgeError::Stream(error)),
194                _ => {}
195            }
196        }
197        Ok(raw_response)
198    }
199}
200
201impl JudgeBuilder {
202    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
203        self.instructions = Some(instructions.into());
204        self
205    }
206
207    pub fn task(mut self, task: impl Into<String>) -> Self {
208        self.task = Some(task.into());
209        self
210    }
211
212    pub fn transcript(mut self, transcript: impl Into<Vec<AgentMessage>>) -> Self {
213        self.context.transcript = Some(transcript.into());
214        self
215    }
216
217    pub fn diff(mut self, diff: impl Into<String>) -> Self {
218        self.context.diff = Some(diff.into());
219        self
220    }
221
222    pub fn file(mut self, path: impl Into<String>, contents: impl Into<String>) -> Self {
223        self.context.files.insert(path.into(), contents.into());
224        self
225    }
226
227    pub fn files<T, U, V>(mut self, files: T) -> Self
228    where
229        T: IntoIterator<Item = (U, V)>,
230        U: Into<String>,
231        V: Into<String>,
232    {
233        self.context.files.extend(files.into_iter().map(|(path, contents)| (path.into(), contents.into())));
234        self
235    }
236
237    pub fn criteria<T, U>(mut self, criteria: T) -> Self
238    where
239        T: IntoIterator<Item = U>,
240        U: Borrow<JudgeCriterionSpec>,
241    {
242        self.criteria = criteria.into_iter().map(|criterion| criterion.borrow().clone()).collect();
243        self
244    }
245
246    pub fn context(mut self, context: JudgeContext) -> Self {
247        self.context = context;
248        self
249    }
250
251    pub fn build(self) -> Result<Judge, JudgeError> {
252        let task = self.task.ok_or_else(|| JudgeError::InvalidInput("judge task must be provided".to_string()))?;
253        let criteria = normalize_criteria(self.criteria)?;
254        let prompt = build_prompt(&self.instructions.unwrap_or_default(), &task, &self.context, &criteria);
255        Ok(Judge { prompt, criteria })
256    }
257}
258
259impl JudgeCriterionSpec {
260    pub fn new(id: impl Into<String>, description: impl Into<String>) -> Self {
261        Self {
262            id: id.into(),
263            description: description.into(),
264            blocking: default_blocking(),
265            weight: default_weight(),
266            threshold: default_threshold(),
267        }
268    }
269
270    pub fn blocking(mut self, blocking: bool) -> Self {
271        self.blocking = blocking;
272        self
273    }
274
275    pub fn weight(mut self, weight: f64) -> Self {
276        self.weight = weight;
277        self
278    }
279
280    pub fn threshold(mut self, threshold: f64) -> Self {
281        self.threshold = threshold;
282        self
283    }
284}
285
286impl JudgeSummary {
287    /// Failure messages for blocking criteria that scored below their threshold.
288    pub fn blocking_failures(&self) -> impl Iterator<Item = String> + '_ {
289        self.criteria
290            .iter()
291            .filter(|criterion| criterion.blocking && !criterion.passed)
292            .map(|criterion| format!("judge criterion `{}`: {}", criterion.id, criterion.reason))
293    }
294}
295
296impl JudgeRubricResponse {
297    pub fn schema() -> Schema {
298        schema_for!(Self)
299    }
300}
301
302fn build_prompt(instructions: &str, task: &str, context: &JudgeContext, criteria: &[JudgeCriterionSpec]) -> String {
303    let mut sections = vec![
304        format!("## Instructions\n\n{instructions}"),
305        format!("## Task\n\nThe agent you're evaluating was given this task: <task>{task}</task>"),
306    ];
307
308    if let Some(transcript) = &context.transcript
309        && !transcript.is_empty()
310    {
311        sections.push(format!(
312            "## Agent Transcript\n\nTranscript of the agent you're evaluating: <transcript>{}</transcript>",
313            format_transcript(transcript)
314        ));
315    }
316
317    if let Some(diff) = &context.diff
318        && !diff.is_empty()
319    {
320        sections.push(format!("## Git diff\n\nGit diff produced by the agent you're evaluating: <diff>{diff}</diff>"));
321    }
322
323    if !context.files.is_empty() {
324        let blocks = context
325            .files
326            .iter()
327            .map(|(path, contents)| format!("<file><path>{path}</path><contents>{contents}</contents></file>"))
328            .collect::<Vec<_>>()
329            .join("\n");
330        sections.push(format!("## File Contents\n\nFiles under evaluation: <files>{blocks}</files>"));
331    }
332
333    let rubric = criteria
334        .iter()
335        .map(|criterion| {
336            format!(
337                "- id: {}\n  blocking: {}\n  weight: {}\n  threshold: {}\n  description: {}",
338                criterion.id, criterion.blocking, criterion.weight, criterion.threshold, criterion.description
339            )
340        })
341        .collect::<Vec<_>>()
342        .join("\n");
343    sections.push(format!("## Rubric criteria\n\n{rubric}"));
344    sections.push(format!(
345        "{}\n{}\n{}\n{}",
346        "Return exactly one result for every criterion ID above and no extra criteria.",
347        "Scores must be normalized numbers from 0.0 to 1.0.",
348        "Respond with ONLY a JSON object matching this schema:",
349        judge_response_schema()
350    ));
351
352    sections.join("\n\n")
353}
354
355fn normalize_criteria(criteria: Vec<JudgeCriterionSpec>) -> Result<Vec<JudgeCriterionSpec>, JudgeError> {
356    if criteria.is_empty() {
357        return Err(JudgeError::InvalidInput("judge criteria must not be empty".to_string()));
358    }
359
360    let mut ids = BTreeSet::new();
361    let mut normalized = Vec::with_capacity(criteria.len());
362    for mut criterion in criteria {
363        criterion.id = criterion.id.trim().to_string();
364        if criterion.id.is_empty() {
365            return Err(JudgeError::InvalidInput("judge criterion id must not be empty".to_string()));
366        }
367        if !ids.insert(criterion.id.clone()) {
368            return Err(JudgeError::InvalidInput(format!("duplicate judge criterion id `{}`", criterion.id)));
369        }
370        if criterion.description.trim().is_empty() {
371            return Err(JudgeError::InvalidInput(format!(
372                "judge criterion `{}` description must not be empty",
373                criterion.id
374            )));
375        }
376        if !criterion.weight.is_finite() || criterion.weight <= 0.0 {
377            return Err(JudgeError::InvalidInput(format!(
378                "judge criterion `{}` weight must be positive and finite",
379                criterion.id
380            )));
381        }
382        if !criterion.threshold.is_finite() || !(0.0..=1.0).contains(&criterion.threshold) {
383            return Err(JudgeError::InvalidInput(format!(
384                "judge criterion `{}` threshold must be between 0.0 and 1.0",
385                criterion.id
386            )));
387        }
388        normalized.push(criterion);
389    }
390    Ok(normalized)
391}
392
393fn extract_json_object(response: &str) -> &str {
394    let trimmed = response.trim();
395    match (trimmed.find('{'), trimmed.rfind('}')) {
396        (Some(start), Some(end)) if start <= end => &trimmed[start..=end],
397        _ => trimmed,
398    }
399}
400
401fn invalid_judgment(reason: String, raw_response: &str) -> JudgeError {
402    JudgeError::InvalidJudgment { reason, raw_response: raw_response.to_string() }
403}
404
405fn judge_response_schema() -> String {
406    serde_json::to_string_pretty(&JudgeRubricResponse::schema()).unwrap()
407}
408
409fn default_blocking() -> bool {
410    true
411}
412
413fn default_weight() -> f64 {
414    1.0
415}
416
417fn default_threshold() -> f64 {
418    1.0
419}
420
421#[cfg(test)]
422mod tests {
423    use super::*;
424    use aether_core::events::AgentMessage;
425    use llm::testing::FakeLlmProvider;
426    use llm::{LlmError, ToolCallRequest};
427
428    const VALID_RESPONSE: &str = r#"{"criteria":[{"id":"behavior","score":1.0,"reason":"correct"},{"id":"clarity","score":0.5,"reason":"brief"}],"overall_reason":"good"}"#;
429
430    #[test]
431    fn judge_builder_builds_prompt_from_context_and_criteria() {
432        let judge = judge()
433            .instructions("be strict")
434            .task("do the thing")
435            .diff("+added line")
436            .file("notes.txt", "beta\n")
437            .criteria([criterion("works", "the task works", true, 2.0, 0.9)])
438            .build()
439            .unwrap();
440
441        assert!(judge.prompt.contains("## Instructions\n\nbe strict"));
442        assert!(judge.prompt.contains("## Task"));
443        assert!(judge.prompt.contains("The agent you're evaluating was given this task: <task>do the thing</task>"));
444        assert!(judge.prompt.contains("## Git diff"));
445        assert!(judge.prompt.contains("Git diff produced by the agent you're evaluating: <diff>+added line</diff>"));
446        assert!(judge.prompt.contains("## File Contents"));
447        assert!(judge.prompt.contains("<path>notes.txt</path>"));
448        assert!(judge.prompt.contains("<contents>beta\n</contents>"));
449        assert!(judge.prompt.contains("## Rubric criteria"));
450        assert!(judge.prompt.contains("blocking: true"));
451        assert!(judge.prompt.contains("threshold: 0.9"));
452        assert!(judge.prompt.contains("weight: 2"));
453        assert!(judge.prompt.contains("Return exactly one result for every criterion ID above and no extra criteria."));
454        assert!(judge.prompt.contains("Respond with ONLY a JSON object matching this schema:"));
455    }
456
457    #[test]
458    fn judge_builder_renders_transcript_context() {
459        let messages = vec![
460            AgentMessage::ToolCall {
461                request: ToolCallRequest {
462                    id: "call_1".to_string(),
463                    name: "bash".to_string(),
464                    arguments: "{}".to_string(),
465                },
466                model_name: "test".to_string(),
467            },
468            AgentMessage::text("msg_1", "all done", true, "test"),
469        ];
470
471        let judge = judge()
472            .task("edit the file")
473            .transcript(messages)
474            .criteria([criterion("behavior", "did it work", true, 1.0, 1.0)])
475            .build()
476            .unwrap();
477
478        assert!(judge.prompt.contains("## Agent Transcript"));
479        assert!(judge.prompt.contains("[tool-call] bash"));
480        assert!(judge.prompt.contains("[agent] all done"));
481    }
482
483    #[test]
484    fn judge_builder_accepts_slice_criteria() {
485        let criteria = vec![criterion("behavior", "does the thing", true, 1.0, 0.8)];
486        let judge = judge().task("do it").criteria(&criteria).build().unwrap();
487        assert_eq!(judge.criteria[0].id, "behavior");
488    }
489
490    #[test]
491    fn judge_summarizes_weighted_rubric() {
492        let judge = judge().task("prompt").criteria(default_criteria()).build().unwrap();
493
494        let summary = judge.summarize(serde_json::from_str(VALID_RESPONSE).unwrap()).unwrap();
495
496        assert!(summary.passed);
497        assert!((summary.score - 0.875).abs() < f64::EPSILON);
498        assert!((summary.criteria[1].score - 0.5).abs() < f64::EPSILON);
499        assert!(summary.reason.contains("all blockers met"));
500    }
501
502    #[test]
503    fn judge_zeroes_score_when_blocker_fails() {
504        let judge = judge().task("prompt").criteria(default_criteria()).build().unwrap();
505        let response = serde_json::from_str(
506            r#"{"criteria":[{"id":"behavior","score":0.75,"reason":"wrong behavior"},{"id":"clarity","score":1.0,"reason":"clear"}],"overall_reason":"bad"}"#,
507        )
508        .unwrap();
509
510        let summary = judge.summarize(response).unwrap();
511
512        assert!(!summary.passed);
513        assert!(summary.score.abs() < f64::EPSILON);
514        assert!(!summary.criteria[0].passed);
515        assert!(summary.reason.contains("one or more blockers failed"));
516    }
517
518    #[test]
519    fn judge_rejects_invalid_criterion_sets() {
520        let judge = judge().task("prompt").criteria(default_criteria()).build().unwrap();
521        for raw_response in [
522            r#"{"criteria":[],"overall_reason":"missing"}"#,
523            r#"{"criteria":[{"id":"behavior","score":1.0,"reason":"ok"},{"id":"behavior","score":1.0,"reason":"ok"}],"overall_reason":"duplicate"}"#,
524            r#"{"criteria":[{"id":"behavior","score":1.0,"reason":"ok"},{"id":"clarity","score":1.0,"reason":"ok"},{"id":"extra","score":1.0,"reason":"ok"}],"overall_reason":"unknown"}"#,
525            r#"{"criteria":[{"id":"behavior","score":1.5,"reason":"bad"},{"id":"clarity","score":1.0,"reason":"ok"}],"overall_reason":"score"}"#,
526        ] {
527            let response = serde_json::from_str(raw_response).unwrap();
528
529            let error = judge.summarize(response).unwrap_err();
530
531            assert!(matches!(error, JudgeError::InvalidJudgment { .. }), "response: {raw_response}");
532        }
533    }
534
535    #[test]
536    fn judge_builder_rejects_invalid_inputs() {
537        for (builder, expected) in [
538            (judge().criteria([criterion("behavior", "ok", true, 1.0, 0.8)]), "judge task must be provided"),
539            (judge().task("prompt"), "judge criteria must not be empty"),
540            (
541                judge().task("prompt").criteria([criterion("", "ok", true, 1.0, 0.8)]),
542                "judge criterion id must not be empty",
543            ),
544            (
545                judge().task("prompt").criteria([criterion("behavior", "ok", true, 0.0, 0.8)]),
546                "weight must be positive and finite",
547            ),
548        ] {
549            let error = builder.build().unwrap_err();
550            assert!(error.to_string().contains(expected), "got: {error}");
551        }
552    }
553
554    #[test]
555    fn blocking_failures_report_only_blocking_criteria_below_threshold() {
556        let criterion = |id: &str, blocking, score: f64| JudgeCriterionSummary {
557            id: id.to_string(),
558            description: "desc".to_string(),
559            blocking,
560            weight: 1.0,
561            threshold: 0.8,
562            score,
563            passed: score >= 0.8,
564            reason: format!("{id} reason"),
565        };
566        let summary = JudgeSummary {
567            passed: false,
568            score: 0.0,
569            reason: "r".to_string(),
570            criteria: vec![
571                criterion("met", true, 0.9),
572                criterion("failed", true, 0.5),
573                criterion("advisory", false, 0.0),
574            ],
575        };
576
577        let failures: Vec<String> = summary.blocking_failures().collect();
578
579        assert_eq!(failures, vec!["judge criterion `failed`: failed reason".to_string()]);
580    }
581
582    #[tokio::test]
583    async fn judge_run_extracts_json_object_from_surrounding_prose() {
584        let response = format!("Here is my assessment:\n{VALID_RESPONSE}");
585        let judge_llm = FakeLlmProvider::with_single_response(vec![LlmResponse::text(&response)]);
586        let judge = judge().task("prompt").criteria(default_criteria()).build().unwrap();
587
588        let summary = judge.run(&judge_llm).await.unwrap();
589
590        assert!(summary.passed);
591    }
592
593    #[tokio::test]
594    async fn judge_run_returns_invalid_json_error_with_raw_response() {
595        let judge_llm = FakeLlmProvider::with_single_response(vec![LlmResponse::text("not json")]);
596        let judge = judge().task("prompt").criteria(default_criteria()).build().unwrap();
597
598        let error = judge.run(&judge_llm).await.unwrap_err();
599
600        let JudgeError::InvalidJson { raw_response, .. } = error else {
601            panic!("expected InvalidJson, got {error:?}");
602        };
603        assert_eq!(raw_response, "not json");
604    }
605
606    #[tokio::test]
607    async fn judge_run_returns_stream_error_on_llm_failure() {
608        let judge_llm = FakeLlmProvider::from_results(vec![vec![Err(LlmError::Other("boom".to_string()))]]);
609        let judge = judge().task("prompt").criteria(default_criteria()).build().unwrap();
610
611        let error = judge.run(&judge_llm).await.unwrap_err();
612
613        assert!(matches!(error, JudgeError::Stream(_)));
614        assert!(error.to_string().contains("boom"));
615    }
616
617    fn criterion(id: &str, description: &str, blocking: bool, weight: f64, threshold: f64) -> JudgeCriterionSpec {
618        JudgeCriterionSpec { id: id.to_string(), description: description.to_string(), blocking, weight, threshold }
619    }
620
621    fn default_criteria() -> Vec<JudgeCriterionSpec> {
622        vec![
623            criterion("behavior", "The behavior is correct.", true, 3.0, 1.0),
624            criterion("clarity", "The response is clear.", false, 1.0, 0.5),
625        ]
626    }
627}
aether_evals/judge.rs

aether_evals/
judge.rs
