ep: Add new prompt format (#52167)

This PR adds `v0318`, which is just like `v0316` but with bigger blocks.
It seems to perform best so far.

It also adds heuristics to avoid placing the marker token on closing
braces, which results in much nicer blocks.

Finally, it fixes a bug with inserting `<|marker_N|>` mid-line in
`v0316` and `v0317`.



Release Notes:

- N/A

crates/edit_prediction_cli/src/format_prompt.rs

@@ -150,6 +150,20 @@ pub fn zeta2_output_for_patch(
         );
     }
 
+    if version == ZetaFormat::V0318SeedMultiRegions {
+        let cursor_in_new = cursor_offset.map(|cursor_offset| {
+            let hunk_start = first_hunk_offset.unwrap_or(0);
+            result.floor_char_boundary((hunk_start + cursor_offset).min(result.len()))
+        });
+        return multi_region::encode_from_old_and_new_v0318(
+            &old_editable_region,
+            &result,
+            cursor_in_new,
+            zeta_prompt::CURSOR_MARKER,
+            multi_region::V0318_END_MARKER,
+        );
+    }
+
     if version == ZetaFormat::V0316SeedMultiRegions {
         let cursor_in_new = cursor_offset.map(|cursor_offset| {
             let hunk_start = first_hunk_offset.unwrap_or(0);

crates/zeta_prompt/src/zeta_prompt.rs

@@ -91,6 +91,8 @@ pub enum ZetaFormat {
     V0306SeedMultiRegions,
     /// Byte-exact marker spans; all intermediate markers emitted; repeated marker means no-edit.
     V0316SeedMultiRegions,
+    /// V0316 with larger block sizes.
+    V0318SeedMultiRegions,
     /// V0316, but marker numbers are relative to the cursor block (e.g. -1, -0, +1).
     V0317SeedMultiRegions,
 }
@@ -242,6 +244,18 @@ pub fn special_tokens_for_format(format: ZetaFormat) -> &'static [&'static str]
             ];
             TOKENS
         }
+        ZetaFormat::V0318SeedMultiRegions => {
+            static TOKENS: &[&str] = &[
+                seed_coder::FIM_SUFFIX,
+                seed_coder::FIM_PREFIX,
+                seed_coder::FIM_MIDDLE,
+                seed_coder::FILE_MARKER,
+                multi_region::V0318_END_MARKER,
+                CURSOR_MARKER,
+                multi_region::MARKER_TAG_PREFIX,
+            ];
+            TOKENS
+        }
         ZetaFormat::V0317SeedMultiRegions => {
             static TOKENS: &[&str] = &[
                 seed_coder::FIM_SUFFIX,
@@ -283,6 +297,7 @@ pub fn token_limits_for_format(format: ZetaFormat) -> (usize, usize) {
         | ZetaFormat::v0226Hashline
         | ZetaFormat::V0306SeedMultiRegions
         | ZetaFormat::V0316SeedMultiRegions
+        | ZetaFormat::V0318SeedMultiRegions
         | ZetaFormat::V0317SeedMultiRegions
         | ZetaFormat::V0304SeedNoEdits => (350, 150),
         ZetaFormat::V0304VariableEdit => (1024, 0),
@@ -303,6 +318,7 @@ pub fn stop_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] {
         | ZetaFormat::V0306SeedMultiRegions
         | ZetaFormat::V0304SeedNoEdits => &[],
         ZetaFormat::V0316SeedMultiRegions => &[multi_region::V0316_END_MARKER],
+        ZetaFormat::V0318SeedMultiRegions => &[multi_region::V0318_END_MARKER],
         ZetaFormat::V0317SeedMultiRegions => &[multi_region::V0317_END_MARKER],
     }
 }
@@ -328,6 +344,7 @@ pub fn excerpt_ranges_for_format(
         | ZetaFormat::V0304SeedNoEdits
         | ZetaFormat::V0306SeedMultiRegions
         | ZetaFormat::V0316SeedMultiRegions
+        | ZetaFormat::V0318SeedMultiRegions
         | ZetaFormat::V0317SeedMultiRegions => (
             ranges.editable_350.clone(),
             ranges.editable_350_context_150.clone(),
@@ -419,6 +436,14 @@ pub fn write_cursor_excerpt_section_for_format(
                 cursor_offset,
             ));
         }
+        ZetaFormat::V0318SeedMultiRegions => {
+            prompt.push_str(&build_v0318_cursor_prefix(
+                path,
+                context,
+                editable_range,
+                cursor_offset,
+            ));
+        }
         ZetaFormat::V0317SeedMultiRegions => {
             prompt.push_str(&build_v0317_cursor_prefix(
                 path,
@@ -486,6 +511,33 @@ fn build_v0316_cursor_prefix(
     section
 }
 
+fn build_v0318_cursor_prefix(
+    path: &Path,
+    context: &str,
+    editable_range: &Range<usize>,
+    cursor_offset: usize,
+) -> String {
+    let mut section = String::new();
+    let path_str = path.to_string_lossy();
+    write!(section, "{}{}\n", seed_coder::FILE_MARKER, path_str).ok();
+
+    section.push_str(&context[..editable_range.start]);
+
+    let editable_text = &context[editable_range.clone()];
+    let cursor_in_editable = cursor_offset - editable_range.start;
+    multi_region::write_editable_with_markers_v0318(
+        &mut section,
+        editable_text,
+        cursor_in_editable,
+        CURSOR_MARKER,
+    );
+
+    if !section.ends_with('\n') {
+        section.push('\n');
+    }
+    section
+}
+
 fn build_v0317_cursor_prefix(
     path: &Path,
     context: &str,
@@ -551,6 +603,7 @@ pub fn format_prompt_with_budget_for_format(
         | ZetaFormat::V0304SeedNoEdits
         | ZetaFormat::V0306SeedMultiRegions
         | ZetaFormat::V0316SeedMultiRegions
+        | ZetaFormat::V0318SeedMultiRegions
         | ZetaFormat::V0317SeedMultiRegions => {
             let mut cursor_section = String::new();
             write_cursor_excerpt_section_for_format(
@@ -649,6 +702,7 @@ pub fn max_edit_event_count_for_format(format: &ZetaFormat) -> usize {
         | ZetaFormat::V0304VariableEdit
         | ZetaFormat::V0306SeedMultiRegions
         | ZetaFormat::V0316SeedMultiRegions
+        | ZetaFormat::V0318SeedMultiRegions
         | ZetaFormat::V0317SeedMultiRegions => 6,
     }
 }
@@ -671,6 +725,7 @@ pub fn get_prefill_for_format(
         ZetaFormat::V0304SeedNoEdits
         | ZetaFormat::V0306SeedMultiRegions
         | ZetaFormat::V0316SeedMultiRegions
+        | ZetaFormat::V0318SeedMultiRegions
         | ZetaFormat::V0317SeedMultiRegions => String::new(),
     }
 }
@@ -684,6 +739,7 @@ pub fn output_end_marker_for_format(format: ZetaFormat) -> Option<&'static str>
         | ZetaFormat::V0304SeedNoEdits
         | ZetaFormat::V0306SeedMultiRegions => Some(seed_coder::END_MARKER),
         ZetaFormat::V0316SeedMultiRegions => Some(multi_region::V0316_END_MARKER),
+        ZetaFormat::V0318SeedMultiRegions => Some(multi_region::V0318_END_MARKER),
         ZetaFormat::V0317SeedMultiRegions => Some(multi_region::V0317_END_MARKER),
         ZetaFormat::V0112MiddleAtEnd
         | ZetaFormat::V0113Ordered
@@ -727,6 +783,22 @@ pub fn encode_patch_as_output_for_format(
                 Ok(None)
             }
         }
+        ZetaFormat::V0318SeedMultiRegions => {
+            let empty_patch = patch.lines().count() <= 3;
+            if empty_patch {
+                let marker_offsets =
+                    multi_region::compute_marker_offsets_v0318(old_editable_region);
+                let marker_num =
+                    multi_region::nearest_marker_number(cursor_offset, &marker_offsets);
+                let tag = multi_region::marker_tag(marker_num);
+                Ok(Some(format!(
+                    "{tag}{tag}{}",
+                    multi_region::V0318_END_MARKER
+                )))
+            } else {
+                Ok(None)
+            }
+        }
         ZetaFormat::V0317SeedMultiRegions => {
             let empty_patch = patch.lines().count() <= 3;
             if empty_patch {
@@ -797,6 +869,10 @@ pub fn parse_zeta2_model_output(
             editable_range_in_context,
             multi_region::apply_marker_span_v0316(old_editable_region, output)?,
         ),
+        ZetaFormat::V0318SeedMultiRegions => (
+            editable_range_in_context,
+            multi_region::apply_marker_span_v0318(old_editable_region, output)?,
+        ),
         ZetaFormat::V0317SeedMultiRegions => (
             editable_range_in_context,
             multi_region::apply_marker_span_v0317(