ep: Make repair parser more robust (#51044)

Oleksiy Syvokon created 1 month ago

Release Notes:

- N/A

Change summary

crates/edit_prediction_cli/src/repair.rs | 21 +++++++++++----------
1 file changed, 11 insertions(+), 10 deletions(-)

Detailed changes

crates/edit_prediction_cli/src/repair.rs

@@ -227,16 +227,17 @@ pub fn needs_repair(example: &Example, confidence_threshold: u8) -> bool {
 /// Handles the `KEEP_PREVIOUS` sentinel by copying the teacher's prediction,
 /// and delegates normal output to `TeacherPrompt::parse`.
 pub fn parse(example: &Example, actual_output: &str) -> Result<(String, Option<ActualCursor>)> {
-    if let Some(last_codeblock) = extract_last_codeblock(actual_output) {
-        if last_codeblock.trim() == KEEP_PREVIOUS {
-            let original = example
-                .predictions
-                .first()
-                .context("no original prediction to keep")?;
-            let patch = original.actual_patch.clone().unwrap_or_default();
-            let cursor = original.actual_cursor.clone();
-            return Ok((patch, cursor));
-        }
+    let last_codeblock =
+        extract_last_codeblock(actual_output).unwrap_or_else(|| actual_output.to_string());
+
+    if last_codeblock.contains(KEEP_PREVIOUS) {
+        let original = example
+            .predictions
+            .first()
+            .context("no original prediction to keep")?;
+        let patch = original.actual_patch.clone().unwrap_or_default();
+        let cursor = original.actual_cursor.clone();
+        return Ok((patch, cursor));
     }
 
     TeacherPrompt::parse(example, actual_output)