diff --git a/crates/edit_prediction/src/cursor_excerpt.rs b/crates/edit_prediction/src/cursor_excerpt.rs index 900d78945ca6ab4fab9c9c60bf13009368c7c77b..18a13130720ca2a5bd313f94437a8c3cfab3c691 100644 --- a/crates/edit_prediction/src/cursor_excerpt.rs +++ b/crates/edit_prediction/src/cursor_excerpt.rs @@ -1,24 +1,15 @@ use language::{BufferSnapshot, Point}; use std::ops::Range; +use text::OffsetRangeExt as _; use zeta_prompt::ExcerptRanges; -/// Pre-computed Point ranges for all editable/context budget combinations. -pub struct ExcerptRangePoints { - pub editable_150: Range, - pub editable_180: Range, - pub editable_350: Range, - pub editable_150_context_350: Range, - pub editable_180_context_350: Range, - pub editable_350_context_150: Range, -} - /// Computes all range variants for a cursor position: editable ranges at 150, 180, and 350 /// token budgets, plus their corresponding context expansions. Returns the full excerpt range /// (union of all context ranges) and the individual sub-ranges as Points. pub fn compute_excerpt_ranges( position: Point, snapshot: &BufferSnapshot, -) -> (Range, ExcerptRangePoints) { +) -> (Range, Range, ExcerptRanges) { let editable_150 = compute_editable_range(snapshot, position, 150); let editable_180 = compute_editable_range(snapshot, position, 180); let editable_350 = compute_editable_range(snapshot, position, 350); @@ -44,37 +35,24 @@ pub fn compute_excerpt_ranges( let full_context = Point::new(full_start_row, 0)..Point::new(full_end_row, snapshot.line_len(full_end_row)); - let ranges = ExcerptRangePoints { - editable_150, - editable_180, - editable_350, - editable_150_context_350, - editable_180_context_350, - editable_350_context_150, - }; - - (full_context, ranges) -} + let full_context_offset_range = full_context.to_offset(snapshot); -/// Converts `ExcerptRangePoints` to byte-offset `ExcerptRanges` relative to `excerpt_start`. -pub fn excerpt_ranges_to_byte_offsets( - ranges: &ExcerptRangePoints, - excerpt_start: usize, - snapshot: &BufferSnapshot, -) -> ExcerptRanges { let to_offset = |range: &Range| -> Range { let start = range.start.to_offset(snapshot); let end = range.end.to_offset(snapshot); - (start - excerpt_start)..(end - excerpt_start) + (start - full_context_offset_range.start)..(end - full_context_offset_range.start) }; - ExcerptRanges { - editable_150: to_offset(&ranges.editable_150), - editable_180: to_offset(&ranges.editable_180), - editable_350: to_offset(&ranges.editable_350), - editable_150_context_350: to_offset(&ranges.editable_150_context_350), - editable_180_context_350: to_offset(&ranges.editable_180_context_350), - editable_350_context_150: to_offset(&ranges.editable_350_context_150), - } + + let ranges = ExcerptRanges { + editable_150: to_offset(&editable_150), + editable_180: to_offset(&editable_180), + editable_350: to_offset(&editable_350), + editable_150_context_350: to_offset(&editable_150_context_350), + editable_180_context_350: to_offset(&editable_180_context_350), + editable_350_context_150: to_offset(&editable_350_context_150), + }; + + (full_context, full_context_offset_range, ranges) } pub fn editable_and_context_ranges_for_cursor_position( diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index cb201870ef636f85e06e0e9cb2def4ec8a149db4..fb6af292fca3b610b5344da146fba558380ad22f 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -2212,9 +2212,7 @@ impl EditPredictionStore { { let http_client = client.http_client(); - let mut token = if let Ok(custom_token) = std::env::var("ZED_PREDICT_EDITS_TOKEN") { - Some(custom_token) - } else if require_auth { + let mut token = if require_auth { Some(llm_token.acquire(&client).await?) } else { llm_token.acquire(&client).await.ok() diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index b68c5df5d306f7db8a66e6be2d6c0847dcfa76dd..658071c9ccfbdf64a9a1ebead7724774cd5cc40e 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -1,4 +1,4 @@ -use crate::cursor_excerpt::{compute_excerpt_ranges, excerpt_ranges_to_byte_offsets}; +use crate::cursor_excerpt::compute_excerpt_ranges; use crate::prediction::EditPredictionResult; use crate::{ CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, @@ -11,7 +11,7 @@ use edit_prediction_types::PredictedCursorPosition; use futures::AsyncReadExt as _; use gpui::{App, AppContext as _, Task, http_client, prelude::*}; use language::language_settings::{OpenAiCompatibleEditPredictionSettings, all_language_settings}; -use language::{BufferSnapshot, OffsetRangeExt as _, ToOffset as _, ToPoint, text_diff}; +use language::{BufferSnapshot, ToOffset as _, ToPoint, text_diff}; use release_channel::AppVersion; use text::{Anchor, Bias}; @@ -24,19 +24,6 @@ use zeta_prompt::{ zeta1::{self, EDITABLE_REGION_END_MARKER}, }; -pub const MAX_CONTEXT_TOKENS: usize = 350; - -pub fn max_editable_tokens(format: ZetaFormat) -> usize { - match format { - ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered => 150, - ZetaFormat::V0114180EditableRegion => 180, - ZetaFormat::V0120GitMergeMarkers => 180, - ZetaFormat::V0131GitMergeMarkersPrefix => 180, - ZetaFormat::V0211Prefill => 180, - ZetaFormat::V0211SeedCoder => 180, - } -} - pub fn request_prediction_with_zeta( store: &mut EditPredictionStore, EditPredictionModelInput { @@ -359,7 +346,8 @@ pub fn zeta2_prompt_input( ) -> (std::ops::Range, zeta_prompt::ZetaPromptInput) { let cursor_point = cursor_offset.to_point(snapshot); - let (full_context, range_points) = compute_excerpt_ranges(cursor_point, snapshot); + let (full_context, full_context_offset_range, excerpt_ranges) = + compute_excerpt_ranges(cursor_point, snapshot); let related_files = crate::filter_redundant_excerpts( related_files, @@ -367,24 +355,17 @@ pub fn zeta2_prompt_input( full_context.start.row..full_context.end.row, ); - let full_context_start_offset = full_context.start.to_offset(snapshot); + let full_context_start_offset = full_context_offset_range.start; let full_context_start_row = full_context.start.row; - let excerpt_ranges = - excerpt_ranges_to_byte_offsets(&range_points, full_context_start_offset, snapshot); - - let editable_range = match preferred_model { - Some(EditPredictionModelKind::Zeta1) => &range_points.editable_350, - _ => match zeta_format { - ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered => &range_points.editable_150, - _ => &range_points.editable_180, - }, + let editable_offset_range = match preferred_model { + Some(EditPredictionModelKind::Zeta1) => excerpt_ranges.editable_350.clone(), + _ => zeta_prompt::excerpt_range_for_format(zeta_format, &excerpt_ranges).0, }; + let absolute_editable_range = full_context_start_offset + editable_offset_range.start + ..full_context_start_offset + editable_offset_range.end; - let editable_offset_range = editable_range.to_offset(snapshot); let cursor_offset_in_excerpt = cursor_offset - full_context_start_offset; - let editable_range_in_excerpt = (editable_offset_range.start - full_context_start_offset) - ..(editable_offset_range.end - full_context_start_offset); let prompt_input = zeta_prompt::ZetaPromptInput { cursor_path: excerpt_path, @@ -392,7 +373,7 @@ pub fn zeta2_prompt_input( .text_for_range(full_context) .collect::() .into(), - editable_range_in_excerpt, + editable_range_in_excerpt: editable_offset_range, cursor_offset_in_excerpt, excerpt_start_row: Some(full_context_start_row), events, @@ -402,7 +383,7 @@ pub fn zeta2_prompt_input( in_open_source_repo: is_open_source, can_collect_data, }; - (editable_offset_range, prompt_input) + (absolute_editable_range, prompt_input) } pub(crate) async fn send_custom_server_request( diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs index fcb33def235fb6953e2da2638c5ee36f094031a3..4a81ac5f670f5ca9dba4797befab4457abc7a98d 100644 --- a/crates/edit_prediction_cli/src/format_prompt.rs +++ b/crates/edit_prediction_cli/src/format_prompt.rs @@ -6,14 +6,14 @@ use crate::{ retrieve_context::run_context_retrieval, }; use anyhow::{Context as _, Result, anyhow}; -use edit_prediction::{cursor_excerpt::editable_and_context_ranges_for_cursor_position, udiff}; +use edit_prediction::{cursor_excerpt::compute_excerpt_ranges, udiff}; use gpui::{AppContext, AsyncApp}; -use language::{Buffer, OffsetRangeExt, Point}; +use language::{Buffer, Point}; use similar::DiffableStr; use std::sync::Arc; use std::{fmt::Write as _, ops::Range}; -use zeta_prompt::ZetaFormat; use zeta_prompt::format_zeta_prompt; +use zeta_prompt::{ZetaFormat, excerpt_range_for_format}; pub async fn run_format_prompt( example: &mut Example, @@ -47,18 +47,15 @@ pub async fn run_format_prompt( let cursor_point = Point::new(prompt_inputs.cursor_row, prompt_inputs.cursor_column); let snapshot = cx.background_spawn(snapshot_fut).await; + let (_, _, excerpt_ranges) = compute_excerpt_ranges(cursor_point, &snapshot); + match args.provider { PredictionProvider::Teacher(_) | PredictionProvider::TeacherNonBatching(_) => { step_progress.set_substatus("formatting teacher prompt"); - let (editable_range, context_range) = editable_and_context_ranges_for_cursor_position( - cursor_point, - &snapshot, - edit_prediction::zeta::max_editable_tokens(ZetaFormat::default()), - edit_prediction::zeta::MAX_CONTEXT_TOKENS, - ); - let editable_range = editable_range.to_offset(&snapshot); - let context_range = context_range.to_offset(&snapshot); + let zeta_format = ZetaFormat::default(); + let (editable_range, context_range) = + excerpt_range_for_format(zeta_format, &excerpt_ranges); let prompt = TeacherPrompt::format_prompt(example, editable_range, context_range); example.prompt = Some(ExamplePrompt { @@ -69,17 +66,11 @@ pub async fn run_format_prompt( provider: args.provider, }); } - PredictionProvider::Zeta2(version) => { + PredictionProvider::Zeta2(zeta_format) => { step_progress.set_substatus("formatting zeta2 prompt"); - let (editable_range, context_range) = editable_and_context_ranges_for_cursor_position( - cursor_point, - &snapshot, - edit_prediction::zeta::max_editable_tokens(version), - edit_prediction::zeta::MAX_CONTEXT_TOKENS, - ); - let editable_range = editable_range.to_offset(&snapshot); - let context_range = context_range.to_offset(&snapshot); + let (editable_range, context_range) = + excerpt_range_for_format(zeta_format, &excerpt_ranges); let context_start = context_range.start; let cursor_offset_in_excerpt = prompt_inputs.cursor_offset - context_start; @@ -93,7 +84,7 @@ pub async fn run_format_prompt( excerpt_start_row: prompt_inputs.excerpt_start_row, events: prompt_inputs.edit_history.clone(), related_files: prompt_inputs.related_files.clone().unwrap_or_default(), - excerpt_ranges: None, + excerpt_ranges: Some(excerpt_ranges), preferred_model: None, in_open_source_repo: example .spec @@ -102,21 +93,24 @@ pub async fn run_format_prompt( .map_or(false, |input| input.in_open_source_repo), can_collect_data: false, }; - let prompt = format_zeta_prompt(&input, version); - let prefill = zeta_prompt::get_prefill(&input, version); + let prompt = format_zeta_prompt(&input, zeta_format); + let prefill = zeta_prompt::get_prefill(&input, zeta_format); let (expected_patch, expected_cursor_offset) = example .spec .expected_patches_with_cursor_positions() .into_iter() .next() .context("expected patches is empty")?; - let expected_output = - zeta2_output_for_patch(&input, &expected_patch, expected_cursor_offset, version)?; - let rejected_output = example - .spec - .rejected_patch - .as_ref() - .and_then(|patch| zeta2_output_for_patch(&input, patch, None, version).ok()); + let expected_output = zeta2_output_for_patch( + &input, + &expected_patch, + expected_cursor_offset, + zeta_format, + )?; + let rejected_output = + example.spec.rejected_patch.as_ref().and_then(|patch| { + zeta2_output_for_patch(&input, patch, None, zeta_format).ok() + }); example.prompt = Some(ExamplePrompt { input: prompt, diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index d8376b5ce312db1166667b89d69bfe9476543a8d..8a2c3477bc2e64fb056f403149e04244897b31cd 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -86,9 +86,9 @@ pub struct ZetaPromptInput { pub enum ZetaFormat { V0112MiddleAtEnd, V0113Ordered, - #[default] V0114180EditableRegion, V0120GitMergeMarkers, + #[default] V0131GitMergeMarkersPrefix, V0211Prefill, V0211SeedCoder, @@ -242,19 +242,11 @@ pub fn clean_zeta2_model_output(output: &str, format: ZetaFormat) -> &str { } } -fn resolve_cursor_region( - input: &ZetaPromptInput, +pub fn excerpt_range_for_format( format: ZetaFormat, -) -> (&str, Range, usize) { - let Some(ranges) = &input.excerpt_ranges else { - return ( - &input.cursor_excerpt, - input.editable_range_in_excerpt.clone(), - input.cursor_offset_in_excerpt, - ); - }; - - let (editable_range, context_range) = match format { + ranges: &ExcerptRanges, +) -> (Range, Range) { + match format { ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered => ( ranges.editable_150.clone(), ranges.editable_150_context_350.clone(), @@ -264,11 +256,25 @@ fn resolve_cursor_region( | ZetaFormat::V0131GitMergeMarkersPrefix | ZetaFormat::V0211Prefill | ZetaFormat::V0211SeedCoder => ( - ranges.editable_180.clone(), - ranges.editable_180_context_350.clone(), + ranges.editable_350.clone(), + ranges.editable_350_context_150.clone(), ), + } +} + +fn resolve_cursor_region( + input: &ZetaPromptInput, + format: ZetaFormat, +) -> (&str, Range, usize) { + let Some(ranges) = &input.excerpt_ranges else { + return ( + &input.cursor_excerpt, + input.editable_range_in_excerpt.clone(), + input.cursor_offset_in_excerpt, + ); }; + let (editable_range, context_range) = excerpt_range_for_format(format, ranges); let context_start = context_range.start; let context_text = &input.cursor_excerpt[context_range]; let adjusted_editable =