//! Quality assessment of predictions using LLM-as-a-judge.//!//! This module uses LLM Batch APIs to evaluate prediction quality.//! Caching is handled by the underlying client. use crate::{ BatchProvider, anthropic_client::AnthropicClient, example::Example, format_prompt::extract_cursor_excerpt_from_example, openai_client::OpenAiClient, parse_output::run_parse_output, paths::LLM_CACHE_DB, progress::{ExampleProgress, Step}, word_diff::unified_to_word_diff,};use anyhow::{Context as _, Result};use serde::{Deserialize, Serialize};use std::sync::OnceLock; /// Arguments for the QA command.#[derive(Debug, Clone, clap::Args)]pub struct QaArgs { /// Use synchronous API instead of batch #[clap(long)] pub no_batch: bool, /// Which LLM provider to use (anthropic or openai) #[clap(long, default_value = "openai")] pub backend: BatchProvider,} fn model_for_backend(backend: BatchProvider) -> &'static str { match backend { BatchProvider::Anthropic => "claude-sonnet-4-5", BatchProvider::Openai => "gpt-5.2", }} /// Result of QA evaluation for a single prediction.#[derive(Debug, Clone, Serialize, Deserialize)]pub struct QaResult { /// Free-form reasoning from the judge. #[serde(default, skip_serializing_if = "Option::is_none")] pub reasoning: Option<String>, /// Does the prediction undo/revert changes the user intentionally made? #[serde(default, skip_serializing_if = "Option::is_none")] pub reverts_edits: Option<bool>, /// Confidence score (1-5) for user acceptance likelihood. #[serde(default, skip_serializing_if = "Option::is_none")] pub confidence: Option<u8>, /// The raw response from the model. #[serde(default, skip_serializing_if = "Option::is_none")] pub response: Option<String>, /// Error message if parsing or request failed. #[serde(default, skip_serializing_if = "Option::is_none")] pub error: Option<String>,} /// Build the assessment prompt for an example.pub fn build_prompt(example: &Example) -> Result<String> { let prediction = example .predictions .first() .context("no predictions available")?; let actual_patch = prediction .actual_patch .as_ref() .context("no actual_patch available (run predict first)")?; let prompt_inputs = example .prompt_inputs .as_ref() .context("prompt_inputs missing (run context retrieval first)")?; let actual_patch_word_diff = unified_to_word_diff(actual_patch); let cursor_excerpt = extract_cursor_excerpt_from_example(example).context("failed to extract cursor excerpt")?; let mut edit_history = String::new(); for event in &prompt_inputs.events { let zeta_prompt::Event::BufferChange { path, old_path, diff, .. } = event.as_ref(); edit_history.push_str(&format!("--- a{}\n", old_path.display())); edit_history.push_str(&format!("+++ b{}\n", path.display())); let diff_word_diff = unified_to_word_diff(&diff); edit_history.push_str(&diff_word_diff); edit_history.push_str("\n\n"); } let prompt_template = crate::prompt_assets::get_prompt("qa.md"); Ok(prompt_template .replace("{edit_history}", &edit_history) .replace("{cursor_excerpt}", &cursor_excerpt) .replace("{actual_patch_word_diff}", &actual_patch_word_diff))} fn extract_codeblock(response: &str) -> Option<String> { let lines: Vec<&str> = response.lines().collect(); for (i, line) in lines.iter().enumerate() { if line.starts_with("```") { let start = i + 1; for (j, end_line) in lines[start..].iter().enumerate() { if end_line.starts_with("```") { return Some(lines[start..start + j].join("\n")); } } return Some(lines[start..].join("\n")); } } None} fn parse_response(response_text: &str) -> QaResult { let codeblock = extract_codeblock(response_text); for text_to_parse in [codeblock.as_deref(), Some(response_text.trim())] { let Some(text) = text_to_parse else { continue; }; if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(text) { return QaResult { reasoning: parsed .get("reasoning") .and_then(|v| v.as_str()) .map(|s| s.to_string()), reverts_edits: parsed.get("reverts_edits").and_then(|v| v.as_bool()), confidence: parsed .get("confidence") .and_then(|v| v.as_u64()) .map(|v| v as u8), response: Some(response_text.to_string()), error: None, }; } } QaResult { reasoning: Some(response_text.to_string()), reverts_edits: None, confidence: None, response: Some(response_text.to_string()), error: Some("Could not parse JSON from response".to_string()), }} static ANTHROPIC_CLIENT_BATCH: OnceLock<AnthropicClient> = OnceLock::new();static ANTHROPIC_CLIENT_PLAIN: OnceLock<AnthropicClient> = OnceLock::new();static OPENAI_CLIENT_BATCH: OnceLock<OpenAiClient> = OnceLock::new();static OPENAI_CLIENT_PLAIN: OnceLock<OpenAiClient> = OnceLock::new(); /// Run QA evaluation for a single example.pub async fn run_qa( example: &mut Example, args: &QaArgs, example_progress: &ExampleProgress,) -> Result<()> { if example .qa .first() .and_then(|q| q.as_ref()) .and_then(|q| q.confidence) .is_some() { return Ok(()); } run_parse_output(example).context("Failed to execute run_parse_output")?; if example.prompt_inputs.is_none() { anyhow::bail!("prompt_inputs missing (run context retrieval first)"); } let step_progress = example_progress.start(Step::Qa); let model = model_for_backend(args.backend); let prompt = build_prompt(example).context("Failed to build QA prompt")?; step_progress.set_substatus("generating"); let response = match args.backend { BatchProvider::Anthropic => { let client = if args.no_batch { ANTHROPIC_CLIENT_PLAIN.get_or_init(|| { AnthropicClient::plain().expect("Failed to create Anthropic client") }) } else { ANTHROPIC_CLIENT_BATCH.get_or_init(|| { AnthropicClient::batch(&LLM_CACHE_DB) .expect("Failed to create Anthropic client") }) }; let messages = vec![anthropic::Message { role: anthropic::Role::User, content: vec![anthropic::RequestContent::Text { text: prompt, cache_control: None, }], }]; let Some(response) = client.generate(model, 1024, messages, None, false).await? else { return Ok(()); }; response .content .iter() .filter_map(|c| match c { anthropic::ResponseContent::Text { text } => Some(text.as_str()), _ => None, }) .collect::<Vec<_>>() .concat() } BatchProvider::Openai => { let client = if args.no_batch { OPENAI_CLIENT_PLAIN .get_or_init(|| OpenAiClient::plain().expect("Failed to create OpenAI client")) } else { OPENAI_CLIENT_BATCH.get_or_init(|| { OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client") }) }; let messages = vec![open_ai::RequestMessage::User { content: open_ai::MessageContent::Plain(prompt), }]; let Some(response) = client.generate(model, 1024, messages, None, false).await? else { return Ok(()); }; response .choices .into_iter() .filter_map(|choice| match choice.message { open_ai::RequestMessage::Assistant { content, .. } => { content.map(|c| match c { open_ai::MessageContent::Plain(text) => text, open_ai::MessageContent::Multipart(parts) => parts .into_iter() .filter_map(|p| match p { open_ai::MessagePart::Text { text } => Some(text), _ => None, }) .collect::<Vec<_>>() .concat(), }) } _ => None, }) .collect::<Vec<_>>() .concat() } }; let result = parse_response(&response); example.qa = example .predictions .iter() .enumerate() .map(|(i, _)| if i == 0 { Some(result.clone()) } else { None }) .collect(); Ok(())} /// Sync batches for QA (upload pending requests, download finished results).pub async fn sync_batches(args: &QaArgs) -> Result<()> { if args.no_batch { return Ok(()); } match args.backend { BatchProvider::Anthropic => { let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| { AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client") }); client.sync_batches().await?; } BatchProvider::Openai => { let client = OPENAI_CLIENT_BATCH.get_or_init(|| { OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client") }); client.sync_batches().await?; } } Ok(())}
//! Quality assessment of predictions using LLM-as-a-judge.
//!
//! This module uses LLM Batch APIs to evaluate prediction quality.
//! Caching is handled by the underlying client.
use crate::{
BatchProvider,
anthropic_client::AnthropicClient,
example::Example,
format_prompt::extract_cursor_excerpt_from_example,
openai_client::OpenAiClient,
parse_output::run_parse_output,
paths::LLM_CACHE_DB,
progress::{ExampleProgress, Step},
word_diff::unified_to_word_diff,
};
use anyhow::{Context as _, Result};
use serde::{Deserialize, Serialize};
use std::sync::OnceLock;
/// Arguments for the QA command.
#[derive(Debug, Clone, clap::Args)]
pub struct QaArgs {
/// Use synchronous API instead of batch
#[clap(long)]
pub no_batch: bool,
/// Which LLM provider to use (anthropic or openai)
#[clap(long, default_value = "openai")]
pub backend: BatchProvider,
}
fn model_for_backend(backend: BatchProvider) -> &'static str {
match backend {
BatchProvider::Anthropic => "claude-sonnet-4-5",
BatchProvider::Openai => "gpt-5.2",
}
}
/// Result of QA evaluation for a single prediction.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QaResult {
/// Free-form reasoning from the judge.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<String>,
/// Does the prediction undo/revert changes the user intentionally made?
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reverts_edits: Option<bool>,
/// Confidence score (1-5) for user acceptance likelihood.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub confidence: Option<u8>,
/// The raw response from the model.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response: Option<String>,
/// Error message if parsing or request failed.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
/// Build the assessment prompt for an example.
pub fn build_prompt(example: &Example) -> Result<String> {
let prediction = example
.predictions
.first()
.context("no predictions available")?;
let actual_patch = prediction
.actual_patch
.as_ref()
.context("no actual_patch available (run predict first)")?;
let prompt_inputs = example
.prompt_inputs
.as_ref()
.context("prompt_inputs missing (run context retrieval first)")?;
let actual_patch_word_diff = unified_to_word_diff(actual_patch);
let cursor_excerpt =
extract_cursor_excerpt_from_example(example).context("failed to extract cursor excerpt")?;
let mut edit_history = String::new();
for event in &prompt_inputs.events {
let zeta_prompt::Event::BufferChange {
path,
old_path,
diff,
..
} = event.as_ref();
edit_history.push_str(&format!("--- a{}\n", old_path.display()));
edit_history.push_str(&format!("+++ b{}\n", path.display()));
let diff_word_diff = unified_to_word_diff(&diff);
edit_history.push_str(&diff_word_diff);
edit_history.push_str("\n\n");
}
let prompt_template = crate::prompt_assets::get_prompt("qa.md");
Ok(prompt_template
.replace("{edit_history}", &edit_history)
.replace("{cursor_excerpt}", &cursor_excerpt)
.replace("{actual_patch_word_diff}", &actual_patch_word_diff))
}
fn extract_codeblock(response: &str) -> Option<String> {
let lines: Vec<&str> = response.lines().collect();
for (i, line) in lines.iter().enumerate() {
if line.starts_with("```") {
let start = i + 1;
for (j, end_line) in lines[start..].iter().enumerate() {
if end_line.starts_with("```") {
return Some(lines[start..start + j].join("\n"));
}
}
return Some(lines[start..].join("\n"));
}
}
None
}
fn parse_response(response_text: &str) -> QaResult {
let codeblock = extract_codeblock(response_text);
for text_to_parse in [codeblock.as_deref(), Some(response_text.trim())] {
let Some(text) = text_to_parse else {
continue;
};
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(text) {
return QaResult {
reasoning: parsed
.get("reasoning")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
reverts_edits: parsed.get("reverts_edits").and_then(|v| v.as_bool()),
confidence: parsed
.get("confidence")
.and_then(|v| v.as_u64())
.map(|v| v as u8),
response: Some(response_text.to_string()),
error: None,
};
}
}
QaResult {
reasoning: Some(response_text.to_string()),
reverts_edits: None,
confidence: None,
response: Some(response_text.to_string()),
error: Some("Could not parse JSON from response".to_string()),
}
}
static ANTHROPIC_CLIENT_BATCH: OnceLock<AnthropicClient> = OnceLock::new();
static ANTHROPIC_CLIENT_PLAIN: OnceLock<AnthropicClient> = OnceLock::new();
static OPENAI_CLIENT_BATCH: OnceLock<OpenAiClient> = OnceLock::new();
static OPENAI_CLIENT_PLAIN: OnceLock<OpenAiClient> = OnceLock::new();
/// Run QA evaluation for a single example.
pub async fn run_qa(
example: &mut Example,
args: &QaArgs,
example_progress: &ExampleProgress,
) -> Result<()> {
if example
.qa
.first()
.and_then(|q| q.as_ref())
.and_then(|q| q.confidence)
.is_some()
{
return Ok(());
}
run_parse_output(example).context("Failed to execute run_parse_output")?;
if example.prompt_inputs.is_none() {
anyhow::bail!("prompt_inputs missing (run context retrieval first)");
}
let step_progress = example_progress.start(Step::Qa);
let model = model_for_backend(args.backend);
let prompt = build_prompt(example).context("Failed to build QA prompt")?;
step_progress.set_substatus("generating");
let response = match args.backend {
BatchProvider::Anthropic => {
let client = if args.no_batch {
ANTHROPIC_CLIENT_PLAIN.get_or_init(|| {
AnthropicClient::plain().expect("Failed to create Anthropic client")
})
} else {
ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
AnthropicClient::batch(&LLM_CACHE_DB)
.expect("Failed to create Anthropic client")
})
};
let messages = vec![anthropic::Message {
role: anthropic::Role::User,
content: vec![anthropic::RequestContent::Text {
text: prompt,
cache_control: None,
}],
}];
let Some(response) = client.generate(model, 1024, messages, None, false).await? else {
return Ok(());
};
response
.content
.iter()
.filter_map(|c| match c {
anthropic::ResponseContent::Text { text } => Some(text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.concat()
}
BatchProvider::Openai => {
let client = if args.no_batch {
OPENAI_CLIENT_PLAIN
.get_or_init(|| OpenAiClient::plain().expect("Failed to create OpenAI client"))
} else {
OPENAI_CLIENT_BATCH.get_or_init(|| {
OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
})
};
let messages = vec![open_ai::RequestMessage::User {
content: open_ai::MessageContent::Plain(prompt),
}];
let Some(response) = client.generate(model, 1024, messages, None, false).await? else {
return Ok(());
};
response
.choices
.into_iter()
.filter_map(|choice| match choice.message {
open_ai::RequestMessage::Assistant { content, .. } => {
content.map(|c| match c {
open_ai::MessageContent::Plain(text) => text,
open_ai::MessageContent::Multipart(parts) => parts
.into_iter()
.filter_map(|p| match p {
open_ai::MessagePart::Text { text } => Some(text),
_ => None,
})
.collect::<Vec<_>>()
.concat(),
})
}
_ => None,
})
.collect::<Vec<_>>()
.concat()
}
};
let result = parse_response(&response);
example.qa = example
.predictions
.iter()
.enumerate()
.map(|(i, _)| if i == 0 { Some(result.clone()) } else { None })
.collect();
Ok(())
}
/// Sync batches for QA (upload pending requests, download finished results).
pub async fn sync_batches(args: &QaArgs) -> Result<()> {
if args.no_batch {
return Ok(());
}
match args.backend {
BatchProvider::Anthropic => {
let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client")
});
client.sync_batches().await?;
}
BatchProvider::Openai => {
let client = OPENAI_CLIENT_BATCH.get_or_init(|| {
OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
});
client.sync_batches().await?;
}
}
Ok(())
}