395 lines
13 KiB
Rust
395 lines
13 KiB
Rust
//! 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;
|
|
use crate::anthropic_client::AnthropicClient;
|
|
use crate::example::Example;
|
|
use crate::format_prompt::extract_cursor_excerpt_from_example;
|
|
use crate::openai_client::OpenAiClient;
|
|
use crate::paths::LLM_CACHE_DB;
|
|
use crate::word_diff::unified_to_word_diff;
|
|
use anyhow::Result;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::io::{BufWriter, Write};
|
|
use std::path::PathBuf;
|
|
|
|
const PROMPT_TEMPLATE: &str = include_str!("prompts/qa.md");
|
|
|
|
/// 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,
|
|
|
|
/// Wait for batch to complete (polls every 30s)
|
|
#[clap(long)]
|
|
pub wait: 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) -> Option<String> {
|
|
let prediction = example.predictions.first()?;
|
|
let actual_patch = prediction.actual_patch.as_ref()?;
|
|
let prompt_inputs = example.prompt_inputs.as_ref()?;
|
|
|
|
let actual_patch_word_diff = unified_to_word_diff(actual_patch);
|
|
|
|
// Format cursor excerpt (reuse from format_prompt)
|
|
let cursor_excerpt = extract_cursor_excerpt_from_example(example)?;
|
|
|
|
let mut edit_history = String::new();
|
|
for event in &prompt_inputs.edit_history {
|
|
match event.as_ref() {
|
|
zeta_prompt::Event::BufferChange {
|
|
path,
|
|
old_path,
|
|
diff,
|
|
predicted: _,
|
|
in_open_source_repo: _,
|
|
} => {
|
|
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");
|
|
}
|
|
}
|
|
}
|
|
|
|
Some(
|
|
PROMPT_TEMPLATE
|
|
.replace("{edit_history}", &edit_history)
|
|
.replace("{cursor_excerpt}", &cursor_excerpt)
|
|
.replace("{actual_patch_word_diff}", &actual_patch_word_diff),
|
|
)
|
|
}
|
|
|
|
/// Extract a code block from a response.
|
|
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
|
|
}
|
|
|
|
/// Parse the LLM response into a QaResult.
|
|
fn parse_response(response_text: &str) -> QaResult {
|
|
let codeblock = extract_codeblock(response_text);
|
|
|
|
// Try parsing codeblock first, then fall back to raw response
|
|
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,
|
|
};
|
|
}
|
|
}
|
|
|
|
// If all parsing attempts fail, return error
|
|
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()),
|
|
}
|
|
}
|
|
|
|
enum QaClient {
|
|
Anthropic(AnthropicClient),
|
|
OpenAi(OpenAiClient),
|
|
}
|
|
|
|
impl QaClient {
|
|
async fn generate(&self, model: &str, max_tokens: u64, prompt: &str) -> Result<Option<String>> {
|
|
match self {
|
|
QaClient::Anthropic(client) => {
|
|
let messages = vec![anthropic::Message {
|
|
role: anthropic::Role::User,
|
|
content: vec![anthropic::RequestContent::Text {
|
|
text: prompt.to_string(),
|
|
cache_control: None,
|
|
}],
|
|
}];
|
|
let response = client
|
|
.generate(model, max_tokens, messages, None, false)
|
|
.await?;
|
|
Ok(response.map(|r| {
|
|
r.content
|
|
.iter()
|
|
.filter_map(|c| match c {
|
|
anthropic::ResponseContent::Text { text } => Some(text.as_str()),
|
|
_ => None,
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("")
|
|
}))
|
|
}
|
|
QaClient::OpenAi(client) => {
|
|
let messages = vec![open_ai::RequestMessage::User {
|
|
content: open_ai::MessageContent::Plain(prompt.to_string()),
|
|
}];
|
|
let response = client
|
|
.generate(model, max_tokens, messages, None, false)
|
|
.await?;
|
|
Ok(response.map(|r| {
|
|
r.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<_>>()
|
|
.join(""),
|
|
})
|
|
}
|
|
_ => None,
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("")
|
|
}))
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn sync_batches(&self) -> Result<()> {
|
|
match self {
|
|
QaClient::Anthropic(client) => client.sync_batches().await,
|
|
QaClient::OpenAi(client) => client.sync_batches().await,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Run the QA evaluation on a set of examples.
|
|
pub async fn run_qa(
|
|
examples: &mut [Example],
|
|
args: &QaArgs,
|
|
output_path: Option<&PathBuf>,
|
|
) -> Result<()> {
|
|
let model = model_for_backend(args.backend);
|
|
let client = match args.backend {
|
|
BatchProvider::Anthropic => {
|
|
if args.no_batch {
|
|
QaClient::Anthropic(AnthropicClient::plain()?)
|
|
} else {
|
|
QaClient::Anthropic(AnthropicClient::batch(&LLM_CACHE_DB)?)
|
|
}
|
|
}
|
|
BatchProvider::Openai => {
|
|
if args.no_batch {
|
|
QaClient::OpenAi(OpenAiClient::plain()?)
|
|
} else {
|
|
QaClient::OpenAi(OpenAiClient::batch(&LLM_CACHE_DB)?)
|
|
}
|
|
}
|
|
};
|
|
|
|
eprintln!(
|
|
"Using model: {}, backend: {:?}, batching: {}",
|
|
model, args.backend, !args.no_batch
|
|
);
|
|
|
|
// First pass: send requests (client handles caching internally)
|
|
let mut prompts: Vec<(usize, String)> = Vec::new();
|
|
let mut skipped_count = 0;
|
|
|
|
for (idx, example) in examples.iter().enumerate() {
|
|
let Some(prompt) = build_prompt(example) else {
|
|
skipped_count += 1;
|
|
continue;
|
|
};
|
|
prompts.push((idx, prompt));
|
|
}
|
|
|
|
if skipped_count > 0 {
|
|
eprintln!("Skipping {} items with missing actual_patch", skipped_count);
|
|
}
|
|
|
|
eprintln!("{} items to process", prompts.len());
|
|
|
|
// Process all items
|
|
let mut results: Vec<(usize, Option<QaResult>)> = Vec::new();
|
|
|
|
if args.no_batch {
|
|
// Synchronous processing
|
|
for (i, (idx, prompt)) in prompts.iter().enumerate() {
|
|
eprint!("\rProcessing {}/{}", i + 1, prompts.len());
|
|
|
|
let response = client.generate(model, 1024, prompt).await?;
|
|
let result = response.map(|text| parse_response(&text));
|
|
results.push((*idx, result));
|
|
}
|
|
eprintln!();
|
|
} else {
|
|
// Queue all for batching
|
|
for (idx, prompt) in &prompts {
|
|
let response = client.generate(model, 1024, prompt).await?;
|
|
let result = response.map(|text| parse_response(&text));
|
|
results.push((*idx, result));
|
|
}
|
|
|
|
// Sync batches (upload pending, download finished)
|
|
client.sync_batches().await?;
|
|
|
|
if args.wait {
|
|
eprintln!("Waiting for batch to complete...");
|
|
loop {
|
|
std::thread::sleep(std::time::Duration::from_secs(30));
|
|
client.sync_batches().await?;
|
|
|
|
// Re-check all items that didn't have results
|
|
let mut all_done = true;
|
|
for (result_idx, (idx, prompt)) in prompts.iter().enumerate() {
|
|
if results[result_idx].1.is_none() {
|
|
let response = client.generate(model, 1024, prompt).await?;
|
|
if let Some(text) = response {
|
|
results[result_idx] = (*idx, Some(parse_response(&text)));
|
|
} else {
|
|
all_done = false;
|
|
}
|
|
}
|
|
}
|
|
|
|
let done_count = results.iter().filter(|(_, r)| r.is_some()).count();
|
|
if all_done {
|
|
break;
|
|
}
|
|
eprintln!("Still waiting... {}/{} results", done_count, prompts.len());
|
|
}
|
|
} else {
|
|
let pending_count = results.iter().filter(|(_, r)| r.is_none()).count();
|
|
if pending_count > 0 {
|
|
eprintln!(
|
|
"Batch submitted. {} pending. Run again later to retrieve results.",
|
|
pending_count
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Build results map by index
|
|
let mut results_by_idx: std::collections::HashMap<usize, QaResult> =
|
|
std::collections::HashMap::new();
|
|
for (idx, result) in results {
|
|
if let Some(r) = result {
|
|
results_by_idx.insert(idx, r);
|
|
}
|
|
}
|
|
|
|
// Output results
|
|
let mut writer: Box<dyn Write> = if let Some(path) = output_path {
|
|
Box::new(BufWriter::new(std::fs::File::create(path)?))
|
|
} else {
|
|
Box::new(std::io::stdout())
|
|
};
|
|
|
|
let mut num_total = 0;
|
|
let mut num_reverts_edits = 0;
|
|
|
|
for (idx, example) in examples.iter_mut().enumerate() {
|
|
// Skip examples that couldn't be processed
|
|
if build_prompt(example).is_none() {
|
|
continue;
|
|
}
|
|
|
|
let result = results_by_idx.get(&idx).cloned();
|
|
|
|
if result.as_ref().and_then(|r| r.reverts_edits) == Some(true) {
|
|
num_reverts_edits += 1;
|
|
}
|
|
num_total += 1;
|
|
|
|
// Populate QA results for each prediction (currently only first prediction is evaluated)
|
|
example.qa = example
|
|
.predictions
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, _)| if i == 0 { result.clone() } else { None })
|
|
.collect();
|
|
|
|
writeln!(writer, "{}", serde_json::to_string(&example)?)?;
|
|
}
|
|
|
|
if let Some(path) = output_path {
|
|
eprintln!("Results written to {}", path.display());
|
|
}
|
|
|
|
eprintln!("Processed: {} items", num_total);
|
|
if num_total > 0 {
|
|
eprintln!(
|
|
"Reverts edits: {} ({:.2}%)",
|
|
num_reverts_edits,
|
|
num_reverts_edits as f64 / num_total as f64 * 100.0
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|