tenant.openagents/omega

1use crate::{
2    PredictionProvider,
3    example::{ActualCursor, Example},
4    format_prompt::{TeacherJumpsPrompt, TeacherPrompt},
5    repair,
6};
7use anyhow::{Context as _, Result};
8use zeta_prompt::{
9    ZetaFormat, parse_zeta2_model_output, parse_zeta2_model_output_as_patch, parsed_output_to_patch,
10};
11
12pub fn run_parse_output(example: &mut Example) -> Result<()> {
13    example
14        .prompt_inputs
15        .as_ref()
16        .context("prompt_inputs required")?;
17
18    let to_parse: Vec<_> = example
19        .predictions
20        .iter()
21        .enumerate()
22        .filter(|(_, p)| !p.actual_output.is_empty())
23        .map(|(ix, p)| (ix, p.actual_output.clone(), p.provider))
24        .collect();
25
26    for (ix, actual_output, provider) in to_parse {
27        let (actual_patch, actual_cursor) =
28            parse_prediction_output(example, &actual_output, provider)?;
29        example.predictions[ix].actual_patch = Some(actual_patch);
30        example.predictions[ix].actual_cursor = actual_cursor;
31    }
32
33    Ok(())
34}
35
36pub fn parse_prediction_output(
37    example: &Example,
38    actual_output: &str,
39    provider: PredictionProvider,
40) -> Result<(String, Option<ActualCursor>)> {
41    match provider {
42        PredictionProvider::Teacher(_, _) | PredictionProvider::TeacherNonBatching(_, _) => {
43            TeacherPrompt::parse(example, actual_output)
44        }
45        PredictionProvider::TeacherJumps(_) | PredictionProvider::TeacherJumpsNonBatching(_) => {
46            TeacherJumpsPrompt::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    if format == ZetaFormat::V0615HashRegions {
68        return Ok((
69            parse_zeta2_model_output_as_patch(actual_output, format, prompt_inputs)?,
70            None,
71        ));
72    }
73
74    let parsed = parse_zeta2_model_output(actual_output, format, prompt_inputs)?;
75    let range_in_excerpt = parsed.range_in_excerpt.clone();
76    let excerpt = prompt_inputs.cursor_excerpt.as_ref();
77    let editable_region_offset = range_in_excerpt.start;
78    let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
79
80    let mut new_text = parsed.new_editable_region.clone();
81    if !new_text.is_empty() && !new_text.ends_with('\n') {
82        new_text.push('\n');
83    }
84
85    let cursor_offset = parsed.cursor_offset_in_new_editable_region;
86    let formatted_diff = parsed_output_to_patch(prompt_inputs, parsed)?;
87
88    let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
89        ActualCursor::from_editable_region(
90            &example.spec.cursor_path,
91            editable_region_cursor_offset,
92            &new_text,
93            excerpt,
94            editable_region_offset,
95            editable_region_start_line,
96        )
97    });
98
99    Ok((formatted_diff, actual_cursor))
100}
101