tokenmiser_quality/
judge.rs

1//! LLM-as-judge pairwise preference scoring.
2
3use anyhow::{anyhow, Result};
4use serde::{Deserialize, Serialize};
5use std::sync::Arc;
6use tokenmiser_providers::{ChatMessage, ChatRequest, ChatResponse, Provider, ProviderRegistry};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
9pub enum JudgeVerdict {
10    /// First response (cheap) wins.
11    A,
12    /// Second response (frontier) wins.
13    B,
14    Tie,
15}
16
17/// Build the judge prompt, presenting the cheap response first. Callers may
18/// invert the order and remap the result to counter positional bias.
19fn judge_prompt(prompt: &str, a: &str, b: &str) -> Vec<ChatMessage> {
20    let system =
21        "You are an impartial judge comparing two assistant responses to the same user prompt. \
22Reply with EXACTLY one token: A, B, or T (for tie). Do not explain."
23            .to_string();
24    let user = format!(
25        "User prompt:\n---\n{prompt}\n---\n\n\
26Response A:\n---\n{a}\n---\n\n\
27Response B:\n---\n{b}\n---\n\n\
28Which response is better? Reply A, B, or T."
29    );
30    vec![
31        ChatMessage {
32            role: "system".into(),
33            content: serde_json::Value::String(system),
34            extra: Default::default(),
35        },
36        ChatMessage {
37            role: "user".into(),
38            content: serde_json::Value::String(user),
39            extra: Default::default(),
40        },
41    ]
42}
43
44/// Run a judge call, resolving `judge_model` through the registry.
45pub async fn judge(
46    registry: &ProviderRegistry,
47    judge_model: &str,
48    user_prompt: &str,
49    cheap_text: &str,
50    frontier_text: &str,
51) -> Result<JudgeVerdict> {
52    let (provider, real) = registry
53        .resolve(judge_model)
54        .map_err(|e| anyhow!("judge resolve: {e}"))?;
55
56    let req = ChatRequest {
57        model: real.clone(),
58        messages: judge_prompt(user_prompt, cheap_text, frontier_text),
59        temperature: Some(0.0),
60        max_tokens: Some(8),
61        top_p: None,
62        stream: None,
63        extra: Default::default(),
64    };
65
66    let resp = provider
67        .complete(&req)
68        .await
69        .map_err(|e| anyhow!("judge call: {e}"))?;
70
71    Ok(parse_verdict(&resp))
72}
73
74fn parse_verdict(resp: &ChatResponse) -> JudgeVerdict {
75    let text = resp
76        .choices
77        .first()
78        .and_then(|c| match &c.message.content {
79            serde_json::Value::String(s) => Some(s.trim().to_uppercase()),
80            _ => None,
81        })
82        .unwrap_or_default();
83
84    // First standalone A/B/T character, skipping ones inside words.
85    for c in text.chars() {
86        match c {
87            'A' => return JudgeVerdict::A,
88            'B' => return JudgeVerdict::B,
89            'T' => return JudgeVerdict::Tie,
90            _ => continue,
91        }
92    }
93    JudgeVerdict::Tie
94}
95
96// Suppress the unused-import warning when this file is read in isolation.
97#[allow(dead_code)]
98fn _provider_typecheck(_: Arc<dyn Provider>) {}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use tokenmiser_providers::{ChatChoice, ChatMessage, Usage};
104
105    fn resp_with(text: &str) -> ChatResponse {
106        ChatResponse {
107            id: "j".into(),
108            object: "chat.completion".into(),
109            created: 0,
110            model: "judge".into(),
111            choices: vec![ChatChoice {
112                index: 0,
113                message: ChatMessage {
114                    role: "assistant".into(),
115                    content: serde_json::Value::String(text.into()),
116                    extra: Default::default(),
117                },
118                finish_reason: Some("stop".into()),
119                logprobs: None,
120            }],
121            usage: Usage::default(),
122            extra: Default::default(),
123        }
124    }
125
126    #[test]
127    fn parses_a_b_t() {
128        assert_eq!(parse_verdict(&resp_with("A")), JudgeVerdict::A);
129        assert_eq!(parse_verdict(&resp_with("B")), JudgeVerdict::B);
130        assert_eq!(parse_verdict(&resp_with("T")), JudgeVerdict::Tie);
131    }
132
133    #[test]
134    fn handles_chatty_judge() {
135        assert_eq!(parse_verdict(&resp_with("A is better")), JudgeVerdict::A);
136        assert_eq!(parse_verdict(&resp_with("I pick B")), JudgeVerdict::B);
137    }
138
139    #[test]
140    fn unparseable_falls_to_tie() {
141        assert_eq!(parse_verdict(&resp_with("hmm")), JudgeVerdict::Tie);
142    }
143}
tokenmiser_quality/judge.rs

tokenmiser_quality/
judge.rs
