1//! Repair predictions that received poor quality signals.
2//!
3//! This module takes examples with predictions, identifies predictions that need
4//! improvement, and uses an LLM to generate improved predictions. It supports
5//! two sources of quality signals:
6//! - QA feedback (reverts_edits or low confidence)
7//! - Computed scores when QA is unavailable (high reversal_ratio or wrong_editable_region)
8
9use crate::{
10 BatchProvider, PredictionProvider,
11 anthropic_client::AnthropicClient,
12 example::{ActualCursor, Example, ExamplePrediction},
13 format_prompt::{TeacherPrompt, extract_last_codeblock},
14 metrics::count_patch_token_changes,
15 openai_client::OpenAiClient,
16 parse_output::run_parse_output,
17 paths::LLM_CACHE_DB,
18 progress::{ExampleProgress, Step},
19 word_diff::unified_to_word_diff,
20};
21use anyhow::{Context as _, Result};
22use std::sync::OnceLock;
23
24const KEEP_PREVIOUS: &str = "KEEP_PREVIOUS";
25
26/// Print a summary report of repair results across all examples.
27pub fn print_report(examples: &[Example], confidence_threshold: u8) {
28 let total = examples.len();
29 let mut no_repair_needed = 0;
30 let mut repaired = 0;
31 let mut repair_failed = 0;
32
33 for example in examples {
34 if !needs_repair(example, confidence_threshold) {
35 no_repair_needed += 1;
36 continue;
37 }
38
39 if has_successful_repair(example) {
40 repaired += 1;
41 } else {
42 repair_failed += 1;
43 }
44 }
45
46 let needed_repair = total - no_repair_needed;
47
48 eprintln!();
49 eprintln!("Repair summary ({total} examples):");
50 eprintln!(
51 " {no_repair_needed}/{total} didn't need repair (confidence > {confidence_threshold})"
52 );
53 if needed_repair > 0 {
54 eprintln!(" {needed_repair}/{total} needed repair:");
55 if repaired > 0 {
56 eprintln!(" {repaired} repaired successfully");
57 }
58 if repair_failed > 0 {
59 eprintln!(" {repair_failed} failed to repair");
60 }
61 }
62}
63
64/// Arguments for the repair command.
65#[derive(Debug, Clone, clap::Args)]
66pub struct RepairArgs {
67 /// Use synchronous API instead of batch
68 #[clap(long)]
69 pub no_batch: bool,
70
71 /// Confidence threshold: repair predictions with confidence <= this value (1-5)
72 #[clap(long, default_value = "2")]
73 pub confidence_threshold: u8,
74
75 /// Which LLM provider to use (anthropic or openai)
76 #[clap(long, default_value = "anthropic")]
77 pub backend: BatchProvider,
78}
79
80fn model_for_backend(backend: BatchProvider) -> &'static str {
81 match backend {
82 BatchProvider::Anthropic => "claude-sonnet-4-6",
83 BatchProvider::Openai => "gpt-5.2",
84 }
85}
86
87/// Build the quality feedback string from QA results.
88fn build_qa_feedback(example: &Example) -> Option<String> {
89 let qa = example.qa.first()?.as_ref()?;
90
91 let qa_reasoning = qa.reasoning.as_deref().unwrap_or("No reasoning provided");
92 let reverts_edits = qa
93 .reverts_edits
94 .map_or("unknown", |v| if v { "yes" } else { "no" });
95 let confidence = qa
96 .confidence
97 .map_or("unknown".to_string(), |v| v.to_string());
98
99 Some(format!(
100 "- **Reverts user edits**: {reverts_edits}\n\
101 - **Confidence score**: {confidence}/5\n\
102 - **Reasoning**: {qa_reasoning}"
103 ))
104}
105
106/// Build the quality feedback string from computed scores when QA is unavailable.
107fn build_score_feedback(example: &Example) -> Option<String> {
108 let score = example.score.first()?;
109
110 let mut issues = Vec::new();
111
112 if score.reversal_ratio > 0.9 {
113 issues.push(format!(
114 "Automated analysis detected a high reversal ratio ({:.2}), which suggests this \
115 prediction may be reverting changes the user intentionally made. Double-check that \
116 the prediction doesn't undo the user's recent edits. If the prediction is actually \
117 fine and the edits are intentional completions rather than reversals, keep it as-is. \
118 If it truly reverts the user's changes, generate an improved prediction that \
119 continues the user's intent instead.",
120 score.reversal_ratio
121 ));
122 }
123
124 if score.wrong_editable_region == Some(true) {
125 issues.push(
126 "Automated analysis detected that the prediction may be modifying code outside \
127 the expected editable region, or producing changes misaligned with the editable \
128 region boundaries. Make sure the prediction only modifies code within the editable \
129 region and is properly aligned."
130 .to_string(),
131 );
132 }
133
134 if issues.is_empty() {
135 return None;
136 }
137
138 let mut feedback = String::from(
139 "No human quality assessment is available, but automated scoring flagged potential issues:\n\n",
140 );
141 for issue in &issues {
142 feedback.push_str(&format!("- {issue}\n"));
143 }
144 feedback.push_str(
145 "\nRemember: if the previous prediction was actually correct, output `KEEP_PREVIOUS`. \
146 If no edits should be made at all and you are unsure how to improve it, output `NO_EDITS`.",
147 );
148
149 Some(feedback)
150}
151
152/// Build the repair message (Turn 3) for a multi-turn conversation.
153///
154/// This message is sent after the original teacher prompt (Turn 1) and
155/// teacher response (Turn 2) to request an improved prediction.
156pub fn build_repair_message(example: &Example) -> Result<String> {
157 let prediction = example
158 .predictions
159 .first()
160 .context("no predictions available")?;
161 let actual_patch = prediction
162 .actual_patch
163 .as_ref()
164 .context("no actual_patch available (run predict first)")?;
165
166 let quality_feedback = build_qa_feedback(example)
167 .or_else(|| build_score_feedback(example))
168 .context("no quality feedback available (need either QA results or computed scores)")?;
169
170 let actual_patch_word_diff = unified_to_word_diff(actual_patch);
171
172 let token_counts = count_patch_token_changes(actual_patch);
173 let mut token_change_info = format!(
174 "\n## Token Change Statistics\n\n\
175 - **Deleted tokens**: {}\n\
176 - **Inserted tokens**: {}",
177 token_counts.deleted_tokens, token_counts.inserted_tokens,
178 );
179 if token_counts.deleted_tokens > 100 || token_counts.inserted_tokens > 100 {
180 token_change_info.push_str(
181 "\n\n> **Note:** The token change count is high. \
182 Consider producing a more scoped edit that targets only the lines \
183 that truly need to change, rather than rewriting large sections.",
184 );
185 }
186
187 let prompt_template = crate::prompt_assets::get_prompt("repair.md");
188 Ok(prompt_template
189 .replace("{actual_patch_word_diff}", &actual_patch_word_diff)
190 .replace("{quality_feedback}", &quality_feedback)
191 .replace("{token_change_info}", &token_change_info))
192}
193
194/// Check if an example needs repair based on QA feedback or computed scores.
195pub fn needs_repair(example: &Example, confidence_threshold: u8) -> bool {
196 // Check QA-based signals first.
197 if let Some(qa) = example.qa.first().and_then(|q| q.as_ref()) {
198 if qa.reverts_edits == Some(true) {
199 return true;
200 }
201
202 if let Some(confidence) = qa.confidence {
203 if confidence <= confidence_threshold {
204 return true;
205 }
206 }
207
208 return false;
209 }
210
211 // When QA is unavailable, fall back to computed score signals.
212 if let Some(score) = example.score.first() {
213 if score.reversal_ratio > 0.9 {
214 return true;
215 }
216
217 if score.wrong_editable_region == Some(true) {
218 return true;
219 }
220 }
221
222 false
223}
224
225/// Parse repair model output into a patch and optional cursor.
226///
227/// Handles the `KEEP_PREVIOUS` sentinel by copying the teacher's prediction,
228/// and delegates normal output to `TeacherPrompt::parse`.
229pub fn parse(example: &Example, actual_output: &str) -> Result<(String, Option<ActualCursor>)> {
230 if let Some(last_codeblock) = extract_last_codeblock(actual_output) {
231 if last_codeblock.trim() == KEEP_PREVIOUS {
232 let original = example
233 .predictions
234 .first()
235 .context("no original prediction to keep")?;
236 let patch = original.actual_patch.clone().unwrap_or_default();
237 let cursor = original.actual_cursor.clone();
238 return Ok((patch, cursor));
239 }
240 }
241
242 TeacherPrompt::parse(example, actual_output)
243}
244
245/// Check if an example already has a successful repair prediction.
246fn has_successful_repair(example: &Example) -> bool {
247 example
248 .predictions
249 .iter()
250 .any(|p| p.provider == PredictionProvider::Repair && p.actual_patch.is_some())
251}
252
253static ANTHROPIC_CLIENT_BATCH: OnceLock<AnthropicClient> = OnceLock::new();
254static ANTHROPIC_CLIENT_PLAIN: OnceLock<AnthropicClient> = OnceLock::new();
255static OPENAI_CLIENT_BATCH: OnceLock<OpenAiClient> = OnceLock::new();
256static OPENAI_CLIENT_PLAIN: OnceLock<OpenAiClient> = OnceLock::new();
257
258/// Run repair for a single example.
259///
260/// This sends a multi-turn conversation to the LLM:
261/// - Turn 1 (User): Original teacher prompt
262/// - Turn 2 (Assistant): Original teacher response
263/// - Turn 3 (User): Repair critique and instructions
264/// - Turn 4 (Assistant): Improved prediction (the response we parse)
265pub async fn run_repair(
266 example: &mut Example,
267 args: &RepairArgs,
268 example_progress: &ExampleProgress,
269) -> Result<()> {
270 if has_successful_repair(example) {
271 return Ok(());
272 }
273
274 if !needs_repair(example, args.confidence_threshold) {
275 return Ok(());
276 }
277
278 run_parse_output(example).context("Failed to execute run_parse_output")?;
279
280 if example.prompt_inputs.is_none() {
281 anyhow::bail!("prompt_inputs missing (run context retrieval first)");
282 }
283
284 if example.predictions.is_empty() {
285 anyhow::bail!("no predictions available (run predict first)");
286 }
287
288 let teacher_prompt = example
289 .prompt
290 .as_ref()
291 .context("prompt missing (run format_prompt first)")?;
292
293 let teacher_response = &example.predictions[0].actual_output;
294 if teacher_response.is_empty() {
295 anyhow::bail!("teacher response is empty (run predict first)");
296 }
297
298 let step_progress = example_progress.start(Step::Repair);
299
300 let model = model_for_backend(args.backend);
301 let repair_message = build_repair_message(example).context("Failed to build repair message")?;
302
303 step_progress.set_substatus("generating");
304
305 let response = match args.backend {
306 BatchProvider::Anthropic => {
307 let client = if args.no_batch {
308 ANTHROPIC_CLIENT_PLAIN.get_or_init(|| {
309 AnthropicClient::plain().expect("Failed to create Anthropic client")
310 })
311 } else {
312 ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
313 AnthropicClient::batch(&LLM_CACHE_DB)
314 .expect("Failed to create Anthropic client")
315 })
316 };
317
318 let messages = vec![
319 // Turn 1: Original teacher prompt
320 anthropic::Message {
321 role: anthropic::Role::User,
322 content: vec![anthropic::RequestContent::Text {
323 text: teacher_prompt.input.clone(),
324 cache_control: None,
325 }],
326 },
327 // Turn 2: Original teacher response
328 anthropic::Message {
329 role: anthropic::Role::Assistant,
330 content: vec![anthropic::RequestContent::Text {
331 text: teacher_response.clone(),
332 cache_control: None,
333 }],
334 },
335 // Turn 3: Repair critique and instructions
336 anthropic::Message {
337 role: anthropic::Role::User,
338 content: vec![anthropic::RequestContent::Text {
339 text: repair_message,
340 cache_control: None,
341 }],
342 },
343 ];
344
345 let Some(response) = client.generate(model, 16384, messages, None, false).await? else {
346 return Ok(());
347 };
348
349 response
350 .content
351 .iter()
352 .filter_map(|c| match c {
353 anthropic::ResponseContent::Text { text } => Some(text.as_str()),
354 _ => None,
355 })
356 .collect::<Vec<_>>()
357 .join("")
358 }
359 BatchProvider::Openai => {
360 let client = if args.no_batch {
361 OPENAI_CLIENT_PLAIN
362 .get_or_init(|| OpenAiClient::plain().expect("Failed to create OpenAI client"))
363 } else {
364 OPENAI_CLIENT_BATCH.get_or_init(|| {
365 OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
366 })
367 };
368
369 let messages = vec![
370 // Turn 1: Original teacher prompt
371 open_ai::RequestMessage::User {
372 content: open_ai::MessageContent::Plain(teacher_prompt.input.clone()),
373 },
374 // Turn 2: Original teacher response
375 open_ai::RequestMessage::Assistant {
376 content: Some(open_ai::MessageContent::Plain(teacher_response.clone())),
377 tool_calls: vec![],
378 },
379 // Turn 3: Repair critique and instructions
380 open_ai::RequestMessage::User {
381 content: open_ai::MessageContent::Plain(repair_message),
382 },
383 ];
384
385 let Some(response) = client.generate(model, 16384, messages, None, false).await? else {
386 return Ok(());
387 };
388
389 response
390 .choices
391 .into_iter()
392 .filter_map(|choice| match choice.message {
393 open_ai::RequestMessage::Assistant { content, .. } => {
394 content.map(|c| match c {
395 open_ai::MessageContent::Plain(text) => text,
396 open_ai::MessageContent::Multipart(parts) => parts
397 .into_iter()
398 .filter_map(|p| match p {
399 open_ai::MessagePart::Text { text } => Some(text),
400 _ => None,
401 })
402 .collect::<Vec<_>>()
403 .join(""),
404 })
405 }
406 _ => None,
407 })
408 .collect::<Vec<_>>()
409 .join("")
410 }
411 };
412
413 let parse_result = parse(example, &response);
414 let err = parse_result
415 .as_ref()
416 .err()
417 .map(|e| format!("Failed to parse repair response: {}", e));
418
419 let (actual_patch, actual_cursor) = parse_result.ok().unzip();
420 let actual_cursor = actual_cursor.flatten();
421
422 example.predictions.push(ExamplePrediction {
423 actual_patch,
424 actual_output: response,
425 actual_cursor,
426 error: err,
427 provider: PredictionProvider::Repair,
428 });
429
430 Ok(())
431}
432
433/// Sync batches for repair (upload pending requests, download finished results).
434pub async fn sync_batches(args: &RepairArgs) -> Result<()> {
435 if args.no_batch {
436 return Ok(());
437 }
438
439 match args.backend {
440 BatchProvider::Anthropic => {
441 let client = ANTHROPIC_CLIENT_BATCH.get_or_init(|| {
442 AnthropicClient::batch(&LLM_CACHE_DB).expect("Failed to create Anthropic client")
443 });
444 client.sync_batches().await?;
445 }
446 BatchProvider::Openai => {
447 let client = OPENAI_CLIENT_BATCH.get_or_init(|| {
448 OpenAiClient::batch(&LLM_CACHE_DB).expect("Failed to create OpenAI client")
449 });
450 client.sync_batches().await?;
451 }
452 }
453
454 Ok(())
455}