zeta: Consolidate logic for picking region sizes, use larger editable region

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<Point>,
-    pub editable_180: Range<Point>,
-    pub editable_350: Range<Point>,
-    pub editable_150_context_350: Range<Point>,
-    pub editable_180_context_350: Range<Point>,
-    pub editable_350_context_150: Range<Point>,
-}
-
 /// 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<Point>, ExcerptRangePoints) {
+) -> (Range<Point>, Range<usize>, 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<Point>| -> Range<usize> {
         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(

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()

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,40 +346,34 @@ pub fn zeta2_prompt_input(
 ) -> (std::ops::Range<usize>, 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,
         excerpt_path.as_ref(),
         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,
         cursor_excerpt: snapshot
             .text_for_range(full_context)
             .collect::<String>()
             .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(

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,

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>, 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<usize>, Range<usize>) {
+    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>, 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 =