diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 5c7ce045121739f341b84dd87d827878550f4048..f6b10ff05c71f9b72c625e7a94518552978357ed 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/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, - cursor_path: &Path, - cursor_row_range: Range, -) -> Vec { - 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." diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs index bf9b43d528db1717f54143e4805e41aefc81f64a..b92c5831239cdaef92fdd2847f30d25215a051fe 100644 --- a/crates/edit_prediction/src/mercury.rs +++ b/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, diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index ccb058e1193eaf2919c286c6e675a907e4af159f..6d7408a3aeb344cd7ef955a1db49c13e561cfa0d 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/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; diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs index f36eaf2799166d6fbd2b7b212003a1a0644b82c4..fe7dff5935aed035e803b1451c8c06df8f79b810 100644 --- a/crates/edit_prediction_cli/src/format_prompt.rs +++ b/crates/edit_prediction_cli/src/format_prompt.rs @@ -95,7 +95,7 @@ pub fn zeta2_output_for_patch( cursor_offset: Option, version: ZetaFormat, ) -> Result { - 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() { diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 2ec12e8bebb4a868c0784e2fe52541a1de580555..2fe353e32250e62c3d3d755540d3efd23c2a627e 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/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) -> Range { + 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, + cursor_path: &Path, + cursor_row_range: Range, +) -> Vec { + 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) { +) -> (&str, Range, Range, 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) }