tenant.openagents/omega

1//! Quality assessment of predictions using LLM-as-a-judge.
2//!
3//! This module uses LLM Batch APIs to evaluate prediction quality.
4//! Caching is handled by the underlying client.
5
6use crate::{
7    BatchProvider,
8    anthropic_client::AnthropicClient,
9    example::Example,
10    format_prompt::extract_cursor_excerpt_from_example,
11    openai_client::OpenAiClient,
12    parse_output::run_parse_output,
13    paths::LLM_CACHE_DB,
14    progress::{ExampleProgress, Step},
15    word_diff::unified_to_word_diff,
16};
17use anyhow::{Context as _, Result};
18use serde::{Deserialize, Serialize};
19use std::sync::OnceLock;
20
21/// Arguments for the QA command.
22#[derive(Debug, Clone, clap::Args)]
23pub struct QaArgs {
24    /// Use synchronous API instead of batch
25    #[clap(long)]
26    pub no_batch: bool,
27
28    /// Which LLM provider to use (anthropic or openai)
29    #[clap(long, default_value = "openai")]
30    pub backend: BatchProvider,
31}
32
33fn model_for_backend(backend: BatchProvider) -> &'static str {
34    match backend {
35        BatchProvider::Anthropic => "claude-sonnet-4-5",
36        BatchProvider::Openai => "gpt-5.2",
37    }
38}
39
40/// Result of QA evaluation for a single prediction.
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct QaResult {
43    /// Free-form reasoning from the judge.
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub reasoning: Option<String>,
46
47    /// Does the prediction undo/revert changes the user intentionally made?
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    pub reverts_edits: Option<bool>,
50
51    /// Confidence score (1-5) for user acceptance likelihood.
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    pub confidence: Option<u8>,
54
55    /// The raw response from the model.
56    #[serde(default, skip_serializing_if = "Option::is_none")]
57    pub response: Option<String>,
58
59    /// Error message if parsing or request failed.
60    #[serde(default, skip_serializing_if = "Option::is_none")]
61    pub error: Option<String>,
62}
63
64/// Build the assessment prompt for an example.
65pub fn build_prompt(example: &Example) -> Result<String> {
66    let prediction = example
67        .predictions
68        .first()
69        .context("no predictions available")?;
70    let actual_patch = prediction
71        .actual_patch
72        .as_ref()
73        .context("no actual_patch available (run predict first)")?;
74    let prompt_inputs = example
75        .prompt_inputs
76        .as_ref()
77        .context("prompt_inputs missing (run context retrieval first)")?;
78
79    let actual_patch_word_diff = unified_to_word_diff(actual_patch);
80
81    let cursor_excerpt =
82        extract_cursor_excerpt_from_example(example).context("failed to extract cursor excerpt")?;
83
84    let mut edit_history = String::new();
85    for event in &prompt_inputs.events {
86        let zeta_prompt::Event::BufferChange {
87            path,
88            old_path,
89            diff,
90            ..
91        } = event.as_ref();
92        edit_history.push_str(&format!("--- a{}\n", old_path.display()));
93        edit_history.push_str(&format!("+++ b{}\n", path.display()));
94        let diff_word_diff = unified_to_word_diff(&diff);
95        edit_history.push_str(&diff_word_diff);
96        edit_history.push_str("\n\n");
97    }
98
99    let prompt_template = crate::prompt_assets::get_prompt("qa.md");
100    Ok(prompt_template
101        .replace("{edit_history}", &edit_history)
102        .replace("{cursor_excerpt}", &cursor_excerpt)
103        .replace("{actual_patch_word_diff}", &actual_patch_word_diff))
104}
105
106fn extract_codeblock(response: &str) -> Option<String> {
107    let lines: Vec<&str> = response.lines().collect();
108    for (i, line) in lines.iter().enumerate() {
109        if line.starts_with("```") {
110            let start = i + 1;
111            for (j, end_line) in lines[start..].iter().enumerate() {
112                if end_line.starts_with("```") {
113                    return Some(lines[start..start + j].join("\n"));
114                }
115            }
116            return Some(lines[start..].join("\n"));
117        }
118    }
119    None
120}
121
122fn parse_response(response_text: &str) -> QaResult {
123    let codeblock = extract_codeblock(response_text);
124
125    for text_to_parse in [codeblock.as_deref(), Some(response_text.trim())] {
126        let Some(text) = text_to_parse else {
127            continue;
128        };
129
130        if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(text) {
131            return QaResult {
132                reasoning: parsed
133                    .get("reasoning")
134                    .and_then(|v| v.as_str())
135                    .map(|s| s.to_string()),
136                reverts_edits: parsed.get("reverts_edits").and_then(|v| v.as_bool()),
137                confidence: parsed
138                    .get("confidence")
139                    .and_then(|v| v.as_u64())
140                    .map(|v| v as u8),
141                response: Some(response_text.to_string()),
142                error: None,
143            };
144        }
145    }
146
147    QaResult {
148        reasoning: Some(response_text.to_string()),
149        reverts_edits: None,
150        confidence: None,
151        response: Some(response_text.to_string()),
152        error: Some("Could not parse JSON from response".to_string()),
153    }
154}
155
156static ANTHROPIC_CLIENT_BATCH: OnceLock<AnthropicClient> = OnceLock::new();
157static ANTHROPIC_CLIENT_PLAIN: OnceLock<AnthropicClient> = OnceLock::new();
158static OPENAI_CLIENT_BATCH: OnceLock<OpenAiClient> = OnceLock::new();
159static OPENAI_CLIENT_PLAIN: OnceLock<OpenAiClient> = OnceLock::new();
160
161/// Run QA evaluation for a single example.
162pub async fn run_qa(
163    example: &mut Example,
164    args: &QaArgs,
165    example_progress: &ExampleProgress,
166) -> Result<()> {
167    if example
168        .qa
169        .first()
170        .and_then(|q| q.as_ref())
171        .and_then(|q| q.confidence)
172        .is_some()
173    {
174        return Ok(());
175    }
176
177    run_parse_output(example).context("Failed to execute run_parse_output")?;
178
179    if example.prompt_inputs.is_none() {
180        anyhow::bail!("prompt_inputs missing (run context retrieval first)");
181    }
182
183    let step_progress = example_progress.start(Step::Qa);
184
185    let model = model_for_backend(args.backend);
186    let prompt = build_prompt(example).context("Failed to build QA prompt")?;
187
188    step_progress.set_substatus("generating");
189
190    let response = match args.backend {
191        BatchProvider::Anthropic => {
192            let client = if args.no_batch {
193                ANTHROPIC_CLIENT_PLAIN.get_or_init(|| {
194                    AnthropicClient::plain().expect("Failed to create Anthropic client")
195                })
196            } else {
197                ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
198                    AnthropicClient::batch(&LLM_CACHE_DB)
199                        .expect("Failed to create Anthropic client")
200                })
201            };
202
203            let messages = vec![anthropic::Message {
204                role: anthropic::Role::User,
205                content: vec![anthropic::RequestContent::Text {
206                    text: prompt,
207                    cache_control: None,
208                }],
209            }];
210
211            let Some(response) = client.generate(model, 1024, messages, None, false).await? else {
212                return Ok(());
213            };
214
215            response
216                .content
217                .iter()
218                .filter_map(|c| match c {
219                    anthropic::ResponseContent::Text { text } => Some(text.as_str()),
220                    _ => None,
221                })
222                .collect::<Vec<_>>()
223                .concat()
224        }
225        BatchProvider::Openai => {
226            let client = if args.no_batch {
227                OPENAI_CLIENT_PLAIN
228                    .get_or_init(|| OpenAiClient::plain().expect("Failed to create OpenAI client"))
229            } else {
230                OPENAI_CLIENT_BATCH.get_or_init(|| {
231                    OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
232                })
233            };
234
235            let messages = vec![open_ai::RequestMessage::User {
236                content: open_ai::MessageContent::Plain(prompt),
237            }];
238
239            let Some(response) = client.generate(model, 1024, messages, None, false).await? else {
240                return Ok(());
241            };
242
243            response
244                .choices
245                .into_iter()
246                .filter_map(|choice| match choice.message {
247                    open_ai::RequestMessage::Assistant { content, .. } => {
248                        content.map(|c| match c {
249                            open_ai::MessageContent::Plain(text) => text,
250                            open_ai::MessageContent::Multipart(parts) => parts
251                                .into_iter()
252                                .filter_map(|p| match p {
253                                    open_ai::MessagePart::Text { text } => Some(text),
254                                    _ => None,
255                                })
256                                .collect::<Vec<_>>()
257                                .concat(),
258                        })
259                    }
260                    _ => None,
261                })
262                .collect::<Vec<_>>()
263                .concat()
264        }
265    };
266
267    let result = parse_response(&response);
268
269    example.qa = example
270        .predictions
271        .iter()
272        .enumerate()
273        .map(|(i, _)| if i == 0 { Some(result.clone()) } else { None })
274        .collect();
275
276    Ok(())
277}
278
279/// Sync batches for QA (upload pending requests, download finished results).
280pub async fn sync_batches(args: &QaArgs) -> Result<()> {
281    if args.no_batch {
282        return Ok(());
283    }
284
285    match args.backend {
286        BatchProvider::Anthropic => {
287            let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
288                AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client")
289            });
290            client.sync_batches().await?;
291        }
292        BatchProvider::Openai => {
293            let client = OPENAI_CLIENT_BATCH.get_or_init(|| {
294                OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
295            });
296            client.sync_batches().await?;
297        }
298    }
299    Ok(())
300}
301