1use crate::{
2 PredictionProvider,
3 example::{ActualCursor, Example},
4 format_prompt::{TeacherMultiRegionPrompt, TeacherPrompt},
5 repair,
6};
7use anyhow::{Context as _, Result};
8use edit_prediction::example_spec::encode_cursor_in_patch;
9use zeta_prompt::{CURSOR_MARKER, ZetaFormat, parse_zeta2_model_output};
10
11pub fn run_parse_output(example: &mut Example) -> Result<()> {
12 example
13 .prompt_inputs
14 .as_ref()
15 .context("prompt_inputs required")?;
16
17 let to_parse: Vec<_> = example
18 .predictions
19 .iter()
20 .enumerate()
21 .filter(|(_, p)| !p.actual_output.is_empty())
22 .map(|(ix, p)| (ix, p.actual_output.clone(), p.provider))
23 .collect();
24
25 for (ix, actual_output, provider) in to_parse {
26 let (actual_patch, actual_cursor) =
27 parse_prediction_output(example, &actual_output, provider)?;
28 example.predictions[ix].actual_patch = Some(actual_patch);
29 example.predictions[ix].actual_cursor = actual_cursor;
30 }
31
32 Ok(())
33}
34
35pub fn parse_prediction_output(
36 example: &Example,
37 actual_output: &str,
38 provider: PredictionProvider,
39) -> Result<(String, Option<ActualCursor>)> {
40 match provider {
41 PredictionProvider::Teacher(_) | PredictionProvider::TeacherNonBatching(_) => {
42 TeacherPrompt::parse(example, actual_output)
43 }
44 PredictionProvider::TeacherMultiRegion(_)
45 | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
46 TeacherMultiRegionPrompt::parse(example, actual_output)
47 }
48 PredictionProvider::Zeta2(version) => parse_zeta2_output(example, actual_output, version),
49 PredictionProvider::Repair => repair::parse(example, actual_output),
50 _ => anyhow::bail!(
51 "parse-output only supports Teacher and Zeta2 providers, got {:?}",
52 provider
53 ),
54 }
55}
56
57fn parse_zeta2_output(
58 example: &Example,
59 actual_output: &str,
60 format: ZetaFormat,
61) -> Result<(String, Option<ActualCursor>)> {
62 let prompt_inputs = example
63 .prompt_inputs
64 .as_ref()
65 .context("prompt_inputs required")?;
66
67 let parsed = parse_zeta2_model_output(actual_output, format, prompt_inputs)?;
68 let range_in_excerpt = parsed.range_in_excerpt;
69
70 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
71 let old_text = excerpt[range_in_excerpt.clone()].to_string();
72 let mut new_text = parsed.new_editable_region;
73
74 let cursor_offset = if let Some(offset) = new_text.find(CURSOR_MARKER) {
75 new_text.replace_range(offset..offset + CURSOR_MARKER.len(), "");
76 Some(offset)
77 } else {
78 None
79 };
80
81 // Normalize trailing newlines for diff generation
82 let mut old_text_normalized = old_text;
83 if !new_text.is_empty() && !new_text.ends_with('\n') {
84 new_text.push('\n');
85 }
86 if !old_text_normalized.is_empty() && !old_text_normalized.ends_with('\n') {
87 old_text_normalized.push('\n');
88 }
89
90 let editable_region_offset = range_in_excerpt.start;
91 let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
92 let editable_region_lines = old_text_normalized.lines().count() as u32;
93
94 let diff = language::unified_diff_with_context(
95 &old_text_normalized,
96 &new_text,
97 editable_region_start_line as u32,
98 editable_region_start_line as u32,
99 editable_region_lines,
100 );
101
102 let formatted_diff = format!(
103 "--- a/{path}\n+++ b/{path}\n{diff}",
104 path = example.spec.cursor_path.to_string_lossy(),
105 );
106
107 let formatted_diff = encode_cursor_in_patch(&formatted_diff, cursor_offset);
108
109 let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
110 ActualCursor::from_editable_region(
111 &example.spec.cursor_path,
112 editable_region_cursor_offset,
113 &new_text,
114 excerpt,
115 editable_region_offset,
116 editable_region_start_line,
117 )
118 });
119
120 Ok((formatted_diff, actual_cursor))
121}