zeta2: Don't remove redundant excerpts on the client (#50886) (#50888)

Ben Kunkle created 1 month ago

Closes #ISSUE

Before you mark this PR as ready for review, make sure that you have:
- [ ] Added a solid test coverage and/or screenshots from doing manual
testing
- [ ] Done a self-review taking into account security and performance
aspects
- [ ] Aligned any UI changes with the [UI
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)

Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/edit_prediction/src/edit_prediction.rs   | 17 -----
crates/edit_prediction/src/mercury.rs           |  2 
crates/edit_prediction/src/zeta.rs              |  6 -
crates/edit_prediction_cli/src/format_prompt.rs |  2 
crates/zeta_prompt/src/zeta_prompt.rs           | 63 ++++++++++++++++--
5 files changed, 57 insertions(+), 33 deletions(-)

Detailed changes

crates/edit_prediction/src/edit_prediction.rs

@@ -2765,23 +2765,6 @@ fn merge_trailing_events_if_needed(
     }
 }
 
-pub(crate) fn filter_redundant_excerpts(
-    mut related_files: Vec<RelatedFile>,
-    cursor_path: &Path,
-    cursor_row_range: Range<u32>,
-) -> Vec<RelatedFile> {
-    for file in &mut related_files {
-        if file.path.as_ref() == cursor_path {
-            file.excerpts.retain(|excerpt| {
-                excerpt.row_range.start < cursor_row_range.start
-                    || excerpt.row_range.end > cursor_row_range.end
-            });
-        }
-    }
-    related_files.retain(|file| !file.excerpts.is_empty());
-    related_files
-}
-
 #[derive(Error, Debug)]
 #[error(
     "You must update to Zed version {minimum_version} or higher to continue using edit predictions."

crates/edit_prediction/src/mercury.rs

@@ -72,7 +72,7 @@ impl Mercury {
                     MAX_REWRITE_TOKENS,
                 );
 
-            let related_files = crate::filter_redundant_excerpts(
+            let related_files = zeta_prompt::filter_redundant_excerpts(
                 related_files,
                 full_path.as_ref(),
                 context_range.start.row..context_range.end.row,

crates/edit_prediction/src/zeta.rs

@@ -397,12 +397,6 @@ pub fn zeta2_prompt_input(
     let (full_context, full_context_offset_range, excerpt_ranges) =
         compute_excerpt_ranges(cursor_point, snapshot);
 
-    let related_files = crate::filter_redundant_excerpts(
-        related_files,
-        excerpt_path.as_ref(),
-        full_context.start.row..full_context.end.row,
-    );
-
     let full_context_start_offset = full_context_offset_range.start;
     let full_context_start_row = full_context.start.row;
 

crates/edit_prediction_cli/src/format_prompt.rs

@@ -95,7 +95,7 @@ pub fn zeta2_output_for_patch(
     cursor_offset: Option<usize>,
     version: ZetaFormat,
 ) -> Result<String> {
-    let (context, editable_range, _) = resolve_cursor_region(input, version);
+    let (context, editable_range, _, _) = resolve_cursor_region(input, version);
     let mut old_editable_region = context[editable_range].to_string();
 
     if !old_editable_region.ends_with_newline() {

crates/zeta_prompt/src/zeta_prompt.rs

@@ -298,14 +298,37 @@ pub fn write_cursor_excerpt_section_for_format(
     }
 }
 
+fn offset_range_to_row_range(text: &str, range: Range<usize>) -> Range<u32> {
+    let start_row = text[0..range.start].matches('\n').count() as u32;
+    let mut end_row = start_row + text[range.clone()].matches('\n').count() as u32;
+    if !text[..range.end].ends_with('\n') {
+        end_row += 1;
+    }
+    return start_row..end_row;
+}
+
 pub fn format_prompt_with_budget_for_format(
     input: &ZetaPromptInput,
     format: ZetaFormat,
     max_tokens: usize,
 ) -> String {
-    let (context, editable_range, cursor_offset) = resolve_cursor_region(input, format);
+    let (context, editable_range, context_range, cursor_offset) =
+        resolve_cursor_region(input, format);
     let path = &*input.cursor_path;
 
+    let related_files = if let Some(cursor_excerpt_start_row) = input.excerpt_start_row {
+        let relative_row_range = offset_range_to_row_range(context, context_range);
+        let row_range = relative_row_range.start + cursor_excerpt_start_row
+            ..relative_row_range.end + cursor_excerpt_start_row;
+        &filter_redundant_excerpts(
+            input.related_files.clone(),
+            input.cursor_path.as_ref(),
+            row_range,
+        )
+    } else {
+        &input.related_files
+    };
+
     match format {
         ZetaFormat::V0211SeedCoder => seed_coder::format_prompt_with_budget(
             path,
@@ -313,7 +336,7 @@ pub fn format_prompt_with_budget_for_format(
             &editable_range,
             cursor_offset,
             &input.events,
-            &input.related_files,
+            &related_files,
             max_tokens,
         ),
         _ => {
@@ -340,7 +363,7 @@ pub fn format_prompt_with_budget_for_format(
             let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens);
 
             let related_files_section = format_related_files_within_budget(
-                &input.related_files,
+                &related_files,
                 "<|file_sep|>",
                 "",
                 budget_after_edit_history,
@@ -355,6 +378,23 @@ pub fn format_prompt_with_budget_for_format(
     }
 }
 
+pub fn filter_redundant_excerpts(
+    mut related_files: Vec<RelatedFile>,
+    cursor_path: &Path,
+    cursor_row_range: Range<u32>,
+) -> Vec<RelatedFile> {
+    for file in &mut related_files {
+        if file.path.as_ref() == cursor_path {
+            file.excerpts.retain(|excerpt| {
+                excerpt.row_range.start < cursor_row_range.start
+                    || excerpt.row_range.end > cursor_row_range.end
+            });
+        }
+    }
+    related_files.retain(|file| !file.excerpts.is_empty());
+    related_files
+}
+
 pub fn get_prefill_for_format(
     format: ZetaFormat,
     context: &str,
@@ -460,19 +500,26 @@ pub fn excerpt_range_for_format(
 pub fn resolve_cursor_region(
     input: &ZetaPromptInput,
     format: ZetaFormat,
-) -> (&str, Range<usize>, usize) {
+) -> (&str, Range<usize>, Range<usize>, usize) {
     let (editable_range, context_range) = excerpt_range_for_format(format, &input.excerpt_ranges);
     let context_start = context_range.start;
-    let context_text = &input.cursor_excerpt[context_range];
+    let context_text = &input.cursor_excerpt[context_range.clone()];
     let adjusted_editable =
         (editable_range.start - context_start)..(editable_range.end - context_start);
     let adjusted_cursor = input.cursor_offset_in_excerpt - context_start;
-
-    (context_text, adjusted_editable, adjusted_cursor)
+    let adjusted_context =
+        (context_range.start - context_start)..(context_range.end - context_start);
+
+    (
+        context_text,
+        adjusted_editable,
+        adjusted_context,
+        adjusted_cursor,
+    )
 }
 
 pub fn get_prefill(input: &ZetaPromptInput, format: ZetaFormat) -> String {
-    let (context, editable_range, _) = resolve_cursor_region(input, format);
+    let (context, editable_range, _, _) = resolve_cursor_region(input, format);
     get_prefill_for_format(format, context, &editable_range)
 }