From 2bd5c218552923ece73e5c7e9afccfa8877d904c Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Mon, 9 Mar 2026 16:58:31 -0700 Subject: [PATCH] zeta: Allow the server to select the editable and context ranges more flexibly (#50975) Release Notes: - N/A --------- Co-authored-by: Ben Kunkle --- Cargo.lock | 1 + crates/codestral/Cargo.toml | 1 + crates/codestral/src/codestral.rs | 45 +- crates/edit_prediction/src/capture_example.rs | 42 +- crates/edit_prediction/src/cursor_excerpt.rs | 517 +++++------------- .../src/edit_prediction_tests.rs | 1 + crates/edit_prediction/src/fim.rs | 36 +- crates/edit_prediction/src/mercury.rs | 65 +-- crates/edit_prediction/src/prediction.rs | 1 + crates/edit_prediction/src/sweep_ai.rs | 1 + crates/edit_prediction/src/zeta.rs | 38 +- .../edit_prediction_cli/src/load_project.rs | 22 +- .../src/retrieve_context.rs | 1 + .../src/reversal_tracking.rs | 1 + crates/zeta_prompt/src/excerpt_ranges.rs | 443 +++++++++++++++ crates/zeta_prompt/src/zeta_prompt.rs | 73 ++- 16 files changed, 761 insertions(+), 527 deletions(-) create mode 100644 crates/zeta_prompt/src/excerpt_ranges.rs diff --git a/Cargo.lock b/Cargo.lock index c549c3b6bfd932bfbec26cebfac3ede79df4d256..3d2ade2ddeab584b8a7ea45590248bdf97e89e57 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3202,6 +3202,7 @@ dependencies = [ "serde", "serde_json", "text", + "zeta_prompt", ] [[package]] diff --git a/crates/codestral/Cargo.toml b/crates/codestral/Cargo.toml index 2addcf110a7c8194538523077d09af9d5104bd0d..0daaee8fb1420c76757ca898655e8dd1a5244d7e 100644 --- a/crates/codestral/Cargo.toml +++ b/crates/codestral/Cargo.toml @@ -22,5 +22,6 @@ log.workspace = true serde.workspace = true serde_json.workspace = true text.workspace = true +zeta_prompt.workspace = true [dev-dependencies] diff --git a/crates/codestral/src/codestral.rs b/crates/codestral/src/codestral.rs index 32436ecc374bef86e3e9a7587acab72741264796..3930e2e873a91618bfae456bc188bbd90ffa64b9 100644 --- a/crates/codestral/src/codestral.rs +++ b/crates/codestral/src/codestral.rs @@ -8,7 +8,7 @@ use gpui::{App, AppContext as _, Context, Entity, Global, SharedString, Task}; use http_client::HttpClient; use icons::IconName; use language::{ - Anchor, Buffer, BufferSnapshot, EditPreview, ToPoint, language_settings::all_language_settings, + Anchor, Buffer, BufferSnapshot, EditPreview, language_settings::all_language_settings, }; use language_model::{ApiKeyState, AuthenticateError, EnvVar, env_var}; use serde::{Deserialize, Serialize}; @@ -18,7 +18,7 @@ use std::{ sync::Arc, time::{Duration, Instant}, }; -use text::{OffsetRangeExt as _, ToOffset}; +use text::ToOffset; pub const CODESTRAL_API_URL: &str = "https://codestral.mistral.ai"; pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(150); @@ -259,28 +259,31 @@ impl EditPredictionDelegate for CodestralEditPredictionDelegate { } let cursor_offset = cursor_position.to_offset(&snapshot); - let cursor_point = cursor_offset.to_point(&snapshot); + const MAX_EDITABLE_TOKENS: usize = 350; const MAX_CONTEXT_TOKENS: usize = 150; - const MAX_REWRITE_TOKENS: usize = 350; - - let (_, context_range) = - cursor_excerpt::editable_and_context_ranges_for_cursor_position( - cursor_point, - &snapshot, - MAX_REWRITE_TOKENS, - MAX_CONTEXT_TOKENS, - ); - - let context_range = context_range.to_offset(&snapshot); - let excerpt_text = snapshot - .text_for_range(context_range.clone()) - .collect::(); - let cursor_within_excerpt = cursor_offset + + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset); + let syntax_ranges = cursor_excerpt::compute_syntax_ranges( + &snapshot, + cursor_offset, + &excerpt_offset_range, + ); + let excerpt_text: String = snapshot.text_for_range(excerpt_point_range).collect(); + let (_, context_range) = zeta_prompt::compute_editable_and_context_ranges( + &excerpt_text, + cursor_offset_in_excerpt, + &syntax_ranges, + MAX_EDITABLE_TOKENS, + MAX_CONTEXT_TOKENS, + ); + let context_text = &excerpt_text[context_range.clone()]; + let cursor_within_excerpt = cursor_offset_in_excerpt .saturating_sub(context_range.start) - .min(excerpt_text.len()); - let prompt = excerpt_text[..cursor_within_excerpt].to_string(); - let suffix = excerpt_text[cursor_within_excerpt..].to_string(); + .min(context_text.len()); + let prompt = context_text[..cursor_within_excerpt].to_string(); + let suffix = context_text[cursor_within_excerpt..].to_string(); let completion_text = match Self::fetch_completion( http_client, diff --git a/crates/edit_prediction/src/capture_example.rs b/crates/edit_prediction/src/capture_example.rs index e0df8cf957747256f86fe5d7f0d63d2ec873d9ca..d21df7868162d279cb18aeea3ef04d4ea9d7be7f 100644 --- a/crates/edit_prediction/src/capture_example.rs +++ b/crates/edit_prediction/src/capture_example.rs @@ -1,12 +1,9 @@ -use crate::{ - StoredEvent, cursor_excerpt::editable_and_context_ranges_for_cursor_position, - example_spec::ExampleSpec, -}; +use crate::{StoredEvent, example_spec::ExampleSpec}; use anyhow::Result; use buffer_diff::BufferDiffSnapshot; use collections::HashMap; use gpui::{App, Entity, Task}; -use language::{Buffer, ToPoint as _}; +use language::Buffer; use project::{Project, WorktreeId}; use std::{collections::hash_map, fmt::Write as _, ops::Range, path::Path, sync::Arc}; use text::{BufferSnapshot as TextBufferSnapshot, Point}; @@ -157,17 +154,34 @@ fn compute_cursor_excerpt( cursor_anchor: language::Anchor, ) -> (String, usize, Range) { use text::ToOffset as _; + use text::ToPoint as _; - let cursor_point = cursor_anchor.to_point(snapshot); - let (_editable_range, context_range) = - editable_and_context_ranges_for_cursor_position(cursor_point, snapshot, 100, 50); - let context_start_offset = context_range.start.to_offset(snapshot); let cursor_offset = cursor_anchor.to_offset(snapshot); - let cursor_offset_in_excerpt = cursor_offset.saturating_sub(context_start_offset); - let excerpt = snapshot - .text_for_range(context_range.clone()) - .collect::(); - (excerpt, cursor_offset_in_excerpt, context_range) + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + crate::cursor_excerpt::compute_cursor_excerpt(snapshot, cursor_offset); + let syntax_ranges = crate::cursor_excerpt::compute_syntax_ranges( + snapshot, + cursor_offset, + &excerpt_offset_range, + ); + let excerpt_text: String = snapshot.text_for_range(excerpt_point_range).collect(); + let (_, context_range) = zeta_prompt::compute_editable_and_context_ranges( + &excerpt_text, + cursor_offset_in_excerpt, + &syntax_ranges, + 100, + 50, + ); + let context_text = excerpt_text[context_range.clone()].to_string(); + let cursor_in_context = cursor_offset_in_excerpt.saturating_sub(context_range.start); + let context_buffer_start = + (excerpt_offset_range.start + context_range.start).to_point(snapshot); + let context_buffer_end = (excerpt_offset_range.start + context_range.end).to_point(snapshot); + ( + context_text, + cursor_in_context, + context_buffer_start..context_buffer_end, + ) } async fn collect_snapshots( diff --git a/crates/edit_prediction/src/cursor_excerpt.rs b/crates/edit_prediction/src/cursor_excerpt.rs index 690e7001bd45ab3d9a995b4dfd43c2e8e297dbe9..27e6cd987a292c71842377226052b665d5a51fbe 100644 --- a/crates/edit_prediction/src/cursor_excerpt.rs +++ b/crates/edit_prediction/src/cursor_excerpt.rs @@ -1,150 +1,30 @@ -use language::{BufferSnapshot, Point}; +use language::{BufferSnapshot, Point, ToPoint as _}; use std::ops::Range; use text::OffsetRangeExt as _; -use zeta_prompt::ExcerptRanges; -/// 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, 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); - let editable_512 = compute_editable_range(snapshot, position, 512); - - let editable_150_context_350 = - expand_context_syntactically_then_linewise(snapshot, editable_150.clone(), 350); - let editable_180_context_350 = - expand_context_syntactically_then_linewise(snapshot, editable_180.clone(), 350); - let editable_350_context_150 = - expand_context_syntactically_then_linewise(snapshot, editable_350.clone(), 150); - let editable_350_context_512 = - expand_context_syntactically_then_linewise(snapshot, editable_350.clone(), 512); - let editable_350_context_1024 = - expand_context_syntactically_then_linewise(snapshot, editable_350.clone(), 1024); - let context_4096 = expand_context_syntactically_then_linewise( - snapshot, - editable_350_context_1024.clone(), - 4096 - 1024, - ); - let context_8192 = - expand_context_syntactically_then_linewise(snapshot, context_4096.clone(), 8192 - 4096); - - let full_start_row = context_8192.start.row; - let full_end_row = context_8192.end.row; - - let full_context = - Point::new(full_start_row, 0)..Point::new(full_end_row, snapshot.line_len(full_end_row)); - - let full_context_offset_range = full_context.to_offset(snapshot); - - let to_offset = |range: &Range| -> Range { - let start = range.start.to_offset(snapshot); - let end = range.end.to_offset(snapshot); - (start - full_context_offset_range.start)..(end - full_context_offset_range.start) - }; - - let ranges = ExcerptRanges { - editable_150: to_offset(&editable_150), - editable_180: to_offset(&editable_180), - editable_350: to_offset(&editable_350), - editable_512: Some(to_offset(&editable_512)), - 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), - editable_350_context_512: Some(to_offset(&editable_350_context_512)), - editable_350_context_1024: Some(to_offset(&editable_350_context_1024)), - context_4096: Some(to_offset(&context_4096)), - context_8192: Some(to_offset(&context_8192)), - }; - - (full_context, full_context_offset_range, ranges) -} - -pub fn editable_and_context_ranges_for_cursor_position( - position: Point, - snapshot: &BufferSnapshot, - editable_region_token_limit: usize, - context_token_limit: usize, -) -> (Range, Range) { - let editable_range = compute_editable_range(snapshot, position, editable_region_token_limit); - - let context_range = expand_context_syntactically_then_linewise( - snapshot, - editable_range.clone(), - context_token_limit, - ); - - (editable_range, context_range) -} +const CURSOR_EXCERPT_TOKEN_BUDGET: usize = 8192; -/// Computes the editable range using a three-phase approach: -/// 1. Expand symmetrically from cursor (75% of budget) -/// 2. Expand to syntax boundaries -/// 3. Continue line-wise in the least-expanded direction -fn compute_editable_range( +/// Computes a cursor excerpt as the largest linewise symmetric region around +/// the cursor that fits within an 8192-token budget. Returns the point range, +/// byte offset range, and the cursor offset relative to the excerpt start. +pub fn compute_cursor_excerpt( snapshot: &BufferSnapshot, - cursor: Point, - token_limit: usize, -) -> Range { - // Phase 1: Expand symmetrically from cursor using 75% of budget. - let initial_budget = (token_limit * 3) / 4; - let (mut start_row, mut end_row, mut remaining_tokens) = - expand_symmetric_from_cursor(snapshot, cursor.row, initial_budget); - - // Add remaining budget from phase 1. - remaining_tokens += token_limit.saturating_sub(initial_budget); - - let original_start = start_row; - let original_end = end_row; - - // Phase 2: Expand to syntax boundaries that fit within budget. - for (boundary_start, boundary_end) in containing_syntax_boundaries(snapshot, start_row, end_row) - { - let tokens_for_start = if boundary_start < start_row { - estimate_tokens_for_rows(snapshot, boundary_start, start_row) - } else { - 0 - }; - let tokens_for_end = if boundary_end > end_row { - estimate_tokens_for_rows(snapshot, end_row + 1, boundary_end + 1) - } else { - 0 - }; - - let total_needed = tokens_for_start + tokens_for_end; - - if total_needed <= remaining_tokens { - if boundary_start < start_row { - start_row = boundary_start; - } - if boundary_end > end_row { - end_row = boundary_end; - } - remaining_tokens = remaining_tokens.saturating_sub(total_needed); - } else { - break; - } - } - - // Phase 3: Continue line-wise in the direction we expanded least during syntax phase. - let expanded_up = original_start.saturating_sub(start_row); - let expanded_down = end_row.saturating_sub(original_end); - - (start_row, end_row, _) = expand_linewise_biased( - snapshot, - start_row, - end_row, - remaining_tokens, - expanded_up <= expanded_down, // prefer_up if we expanded less upward - ); - - let start = Point::new(start_row, 0); - let end = Point::new(end_row, snapshot.line_len(end_row)); - start..end + cursor_offset: usize, +) -> (Range, Range, usize) { + let cursor_point = cursor_offset.to_point(snapshot); + let cursor_row = cursor_point.row; + let (start_row, end_row, _) = + expand_symmetric_from_cursor(snapshot, cursor_row, CURSOR_EXCERPT_TOKEN_BUDGET); + + let excerpt_range = Point::new(start_row, 0)..Point::new(end_row, snapshot.line_len(end_row)); + let excerpt_offset_range = excerpt_range.to_offset(snapshot); + let cursor_offset_in_excerpt = cursor_offset - excerpt_offset_range.start; + + ( + excerpt_range, + excerpt_offset_range, + cursor_offset_in_excerpt, + ) } /// Expands symmetrically from cursor, one line at a time, alternating down then up. @@ -157,7 +37,6 @@ fn expand_symmetric_from_cursor( let mut start_row = cursor_row; let mut end_row = cursor_row; - // Account for the cursor's line. let cursor_line_tokens = line_token_count(snapshot, cursor_row); token_budget = token_budget.saturating_sub(cursor_line_tokens); @@ -169,7 +48,6 @@ fn expand_symmetric_from_cursor( break; } - // Expand down first (slight forward bias for edit prediction). if can_expand_down { let next_row = end_row + 1; let line_tokens = line_token_count(snapshot, next_row); @@ -181,7 +59,6 @@ fn expand_symmetric_from_cursor( } } - // Then expand up. if can_expand_up && token_budget > 0 { let next_row = start_row - 1; let line_tokens = line_token_count(snapshot, next_row); @@ -197,74 +74,6 @@ fn expand_symmetric_from_cursor( (start_row, end_row, token_budget) } -/// Expands line-wise with a bias toward one direction. -/// Returns (start_row, end_row, remaining_tokens). -fn expand_linewise_biased( - snapshot: &BufferSnapshot, - mut start_row: u32, - mut end_row: u32, - mut remaining_tokens: usize, - prefer_up: bool, -) -> (u32, u32, usize) { - loop { - let can_expand_up = start_row > 0; - let can_expand_down = end_row < snapshot.max_point().row; - - if remaining_tokens == 0 || (!can_expand_up && !can_expand_down) { - break; - } - - let mut expanded = false; - - // Try preferred direction first. - if prefer_up { - if can_expand_up { - let next_row = start_row - 1; - let line_tokens = line_token_count(snapshot, next_row); - if line_tokens <= remaining_tokens { - start_row = next_row; - remaining_tokens = remaining_tokens.saturating_sub(line_tokens); - expanded = true; - } - } - if can_expand_down && remaining_tokens > 0 { - let next_row = end_row + 1; - let line_tokens = line_token_count(snapshot, next_row); - if line_tokens <= remaining_tokens { - end_row = next_row; - remaining_tokens = remaining_tokens.saturating_sub(line_tokens); - expanded = true; - } - } - } else { - if can_expand_down { - let next_row = end_row + 1; - let line_tokens = line_token_count(snapshot, next_row); - if line_tokens <= remaining_tokens { - end_row = next_row; - remaining_tokens = remaining_tokens.saturating_sub(line_tokens); - expanded = true; - } - } - if can_expand_up && remaining_tokens > 0 { - let next_row = start_row - 1; - let line_tokens = line_token_count(snapshot, next_row); - if line_tokens <= remaining_tokens { - start_row = next_row; - remaining_tokens = remaining_tokens.saturating_sub(line_tokens); - expanded = true; - } - } - } - - if !expanded { - break; - } - } - - (start_row, end_row, remaining_tokens) -} - /// Typical number of string bytes per token for the purposes of limiting model input. This is /// intentionally low to err on the side of underestimating limits. pub(crate) const BYTES_PER_TOKEN_GUESS: usize = 3; @@ -277,113 +86,50 @@ fn line_token_count(snapshot: &BufferSnapshot, row: u32) -> usize { guess_token_count(snapshot.line_len(row) as usize).max(1) } -/// Estimates token count for rows in range [start_row, end_row). -fn estimate_tokens_for_rows(snapshot: &BufferSnapshot, start_row: u32, end_row: u32) -> usize { - let mut tokens = 0; - for row in start_row..end_row { - tokens += line_token_count(snapshot, row); - } - tokens -} - -/// Returns an iterator of (start_row, end_row) for successively larger syntax nodes -/// containing the given row range. Smallest containing node first. -fn containing_syntax_boundaries( +/// Computes the byte offset ranges of all syntax nodes containing the cursor, +/// ordered from innermost to outermost. The offsets are relative to +/// `excerpt_offset_range.start`. +pub fn compute_syntax_ranges( snapshot: &BufferSnapshot, - start_row: u32, - end_row: u32, -) -> impl Iterator { - let range = Point::new(start_row, 0)..Point::new(end_row, snapshot.line_len(end_row)); + cursor_offset: usize, + excerpt_offset_range: &Range, +) -> Vec> { + let cursor_point = cursor_offset.to_point(snapshot); + let range = cursor_point..cursor_point; let mut current = snapshot.syntax_ancestor(range); - let mut last_rows: Option<(u32, u32)> = None; - - std::iter::from_fn(move || { - while let Some(node) = current.take() { - let node_start_row = node.start_position().row as u32; - let node_end_row = node.end_position().row as u32; - let rows = (node_start_row, node_end_row); - - current = node.parent(); - - // Skip nodes that don't extend beyond our range. - if node_start_row >= start_row && node_end_row <= end_row { - continue; - } + let mut ranges = Vec::new(); + let mut last_range: Option<(usize, usize)> = None; - // Skip if same as last returned (some nodes have same span). - if last_rows == Some(rows) { - continue; - } + while let Some(node) = current.take() { + let node_start = node.start_byte(); + let node_end = node.end_byte(); + let key = (node_start, node_end); - last_rows = Some(rows); - return Some(rows); - } - None - }) -} + current = node.parent(); -/// Expands context by first trying to reach syntax boundaries, -/// then expanding line-wise only if no syntax expansion occurred. -fn expand_context_syntactically_then_linewise( - snapshot: &BufferSnapshot, - editable_range: Range, - context_token_limit: usize, -) -> Range { - let mut start_row = editable_range.start.row; - let mut end_row = editable_range.end.row; - let mut remaining_tokens = context_token_limit; - let mut did_syntax_expand = false; - - // Phase 1: Try to expand to containing syntax boundaries, picking the largest that fits. - for (boundary_start, boundary_end) in containing_syntax_boundaries(snapshot, start_row, end_row) - { - let tokens_for_start = if boundary_start < start_row { - estimate_tokens_for_rows(snapshot, boundary_start, start_row) - } else { - 0 - }; - let tokens_for_end = if boundary_end > end_row { - estimate_tokens_for_rows(snapshot, end_row + 1, boundary_end + 1) - } else { - 0 - }; - - let total_needed = tokens_for_start + tokens_for_end; - - if total_needed <= remaining_tokens { - if boundary_start < start_row { - start_row = boundary_start; - } - if boundary_end > end_row { - end_row = boundary_end; - } - remaining_tokens = remaining_tokens.saturating_sub(total_needed); - did_syntax_expand = true; - } else { - break; + if last_range == Some(key) { + continue; } - } + last_range = Some(key); - // Phase 2: Only expand line-wise if no syntax expansion occurred. - if !did_syntax_expand { - (start_row, end_row, _) = - expand_linewise_biased(snapshot, start_row, end_row, remaining_tokens, true); + let start = node_start.saturating_sub(excerpt_offset_range.start); + let end = node_end + .min(excerpt_offset_range.end) + .saturating_sub(excerpt_offset_range.start); + ranges.push(start..end); } - let start = Point::new(start_row, 0); - let end = Point::new(end_row, snapshot.line_len(end_row)); - start..end + ranges } -use language::ToOffset as _; - #[cfg(test)] mod tests { use super::*; - use gpui::{App, AppContext}; + use gpui::{App, AppContext as _}; use indoc::indoc; use language::{Buffer, rust_lang}; use util::test::{TextRangeMarker, marked_text_ranges_by}; + use zeta_prompt::compute_editable_and_context_ranges; struct TestCase { name: &'static str, @@ -400,7 +146,18 @@ mod tests { // [ ] = expected context range let test_cases = vec![ TestCase { - name: "cursor near end of function - expands to syntax boundaries", + name: "small function fits entirely in editable and context", + marked_text: indoc! {r#" + [«fn foo() { + let x = 1;ˇ + let y = 2; + }»] + "#}, + editable_token_limit: 30, + context_token_limit: 60, + }, + TestCase { + name: "cursor near end of function - editable expands to syntax boundaries", marked_text: indoc! {r#" [fn first() { let a = 1; @@ -413,12 +170,11 @@ mod tests { println!("{}", x + y);ˇ }»] "#}, - // 18 tokens - expands symmetrically then to syntax boundaries editable_token_limit: 18, context_token_limit: 35, }, TestCase { - name: "cursor at function start - expands to syntax boundaries", + name: "cursor at function start - editable expands to syntax boundaries", marked_text: indoc! {r#" [fn before() { « let a = 1; @@ -434,12 +190,11 @@ mod tests { let b = 2; }] "#}, - // 25 tokens - expands symmetrically then to syntax boundaries editable_token_limit: 25, context_token_limit: 50, }, TestCase { - name: "tiny budget - just lines around cursor", + name: "tiny budget - just lines around cursor, no syntax expansion", marked_text: indoc! {r#" fn outer() { [ let line1 = 1; @@ -451,22 +206,9 @@ mod tests { let line7 = 7; } "#}, - // 12 tokens (~36 bytes) = just the cursor line with tiny budget editable_token_limit: 12, context_token_limit: 24, }, - TestCase { - name: "small function fits entirely", - marked_text: indoc! {r#" - [«fn foo() { - let x = 1;ˇ - let y = 2; - }»] - "#}, - // Plenty of budget for this small function - editable_token_limit: 30, - context_token_limit: 60, - }, TestCase { name: "context extends beyond editable", marked_text: indoc! {r#" @@ -476,13 +218,11 @@ mod tests { fn fourth() { let d = 4; }» fn fifth() { let e = 5; }] "#}, - // Small editable, larger context editable_token_limit: 25, context_token_limit: 45, }, - // Tests for syntax-aware editable and context expansion TestCase { - name: "cursor in first if-statement - expands to syntax boundaries", + name: "cursor in first if-block - editable expands to syntax boundaries", marked_text: indoc! {r#" [«fn before() { } @@ -503,13 +243,11 @@ mod tests { fn after() { }] "#}, - // 35 tokens allows expansion to include function header and first two if blocks editable_token_limit: 35, - // 60 tokens allows context to include the whole file context_token_limit: 60, }, TestCase { - name: "cursor in middle if-statement - expands to syntax boundaries", + name: "cursor in middle if-block - editable spans surrounding blocks", marked_text: indoc! {r#" [fn before() { } @@ -530,13 +268,11 @@ mod tests { fn after() { }] "#}, - // 40 tokens allows expansion to surrounding if blocks editable_token_limit: 40, - // 60 tokens allows context to include the whole file context_token_limit: 60, }, TestCase { - name: "cursor near bottom of long function - editable expands toward syntax, context reaches function", + name: "cursor near bottom of long function - context reaches function boundary", marked_text: indoc! {r#" [fn other() { } @@ -556,11 +292,30 @@ mod tests { fn another() { }»] "#}, - // 40 tokens for editable - allows several lines plus syntax expansion editable_token_limit: 40, - // 55 tokens - enough for function but not whole file context_token_limit: 55, }, + TestCase { + name: "zero context budget - context equals editable", + marked_text: indoc! {r#" + fn before() { + let p = 1; + let q = 2; + [«} + + fn foo() { + let x = 1;ˇ + let y = 2; + } + »] + fn after() { + let r = 3; + let s = 4; + } + "#}, + editable_token_limit: 15, + context_token_limit: 0, + }, ]; for test_case in test_cases { @@ -580,75 +335,63 @@ mod tests { let cursor_ranges = ranges.remove(&cursor_marker).unwrap_or_default(); let expected_editable = ranges.remove(&editable_marker).unwrap_or_default(); let expected_context = ranges.remove(&context_marker).unwrap_or_default(); - assert_eq!(expected_editable.len(), 1); - assert_eq!(expected_context.len(), 1); + assert_eq!(expected_editable.len(), 1, "{}", test_case.name); + assert_eq!(expected_context.len(), 1, "{}", test_case.name); - cx.new(|cx| { + cx.new(|cx: &mut gpui::Context| { let text = text.trim_end_matches('\n'); let buffer = Buffer::local(text, cx).with_language(rust_lang(), cx); let snapshot = buffer.snapshot(); let cursor_offset = cursor_ranges[0].start; - let cursor_point = snapshot.offset_to_point(cursor_offset); - let expected_editable_start = snapshot.offset_to_point(expected_editable[0].start); - let expected_editable_end = snapshot.offset_to_point(expected_editable[0].end); - let expected_context_start = snapshot.offset_to_point(expected_context[0].start); - let expected_context_end = snapshot.offset_to_point(expected_context[0].end); - - let (actual_editable, actual_context) = - editable_and_context_ranges_for_cursor_position( - cursor_point, - &snapshot, - test_case.editable_token_limit, - test_case.context_token_limit, - ); - - let range_text = |start: Point, end: Point| -> String { - snapshot.text_for_range(start..end).collect() + + let (_, excerpt_offset_range, cursor_offset_in_excerpt) = + compute_cursor_excerpt(&snapshot, cursor_offset); + let excerpt_text: String = snapshot + .text_for_range(excerpt_offset_range.clone()) + .collect(); + let syntax_ranges = + compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range); + + let (actual_editable, actual_context) = compute_editable_and_context_ranges( + &excerpt_text, + cursor_offset_in_excerpt, + &syntax_ranges, + test_case.editable_token_limit, + test_case.context_token_limit, + ); + + let to_buffer_range = |range: Range| -> Range { + (excerpt_offset_range.start + range.start) + ..(excerpt_offset_range.start + range.end) }; - let editable_match = actual_editable.start == expected_editable_start - && actual_editable.end == expected_editable_end; - let context_match = actual_context.start == expected_context_start - && actual_context.end == expected_context_end; + let actual_editable = to_buffer_range(actual_editable); + let actual_context = to_buffer_range(actual_context); + + let expected_editable_range = expected_editable[0].clone(); + let expected_context_range = expected_context[0].clone(); + + let editable_match = actual_editable == expected_editable_range; + let context_match = actual_context == expected_context_range; if !editable_match || !context_match { + let range_text = |range: &Range| { + snapshot.text_for_range(range.clone()).collect::() + }; + println!("\n=== FAILED: {} ===", test_case.name); if !editable_match { - println!( - "\nExpected editable ({:?}..{:?}):", - expected_editable_start, expected_editable_end - ); - println!( - "---\n{}---", - range_text(expected_editable_start, expected_editable_end) - ); - println!( - "\nActual editable ({:?}..{:?}):", - actual_editable.start, actual_editable.end - ); - println!( - "---\n{}---", - range_text(actual_editable.start, actual_editable.end) - ); + println!("\nExpected editable ({:?}):", expected_editable_range); + println!("---\n{}---", range_text(&expected_editable_range)); + println!("\nActual editable ({:?}):", actual_editable); + println!("---\n{}---", range_text(&actual_editable)); } if !context_match { - println!( - "\nExpected context ({:?}..{:?}):", - expected_context_start, expected_context_end - ); - println!( - "---\n{}---", - range_text(expected_context_start, expected_context_end) - ); - println!( - "\nActual context ({:?}..{:?}):", - actual_context.start, actual_context.end - ); - println!( - "---\n{}---", - range_text(actual_context.start, actual_context.end) - ); + println!("\nExpected context ({:?}):", expected_context_range); + println!("---\n{}---", range_text(&expected_context_range)); + println!("\nActual context ({:?}):", actual_context); + println!("---\n{}---", range_text(&actual_context)); } panic!("Test '{}' failed - see output above", test_case.name); } diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 1ff77fd900db80894b973e79d8fe69e9d65a1e3b..66d9c940dda21d7068ad6dec0976520dee2750e7 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -1890,6 +1890,7 @@ async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) { cursor_offset_in_excerpt: 0, excerpt_start_row: None, excerpt_ranges: Default::default(), + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, diff --git a/crates/edit_prediction/src/fim.rs b/crates/edit_prediction/src/fim.rs index 79df739e60bc28ba5c6b9f53699dcf398fc8310e..1a64506f00285791a83c38943253157137d592f1 100644 --- a/crates/edit_prediction/src/fim.rs +++ b/crates/edit_prediction/src/fim.rs @@ -6,12 +6,12 @@ use crate::{ use anyhow::{Context as _, Result, anyhow}; use gpui::{App, AppContext as _, Entity, Task}; use language::{ - Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, ToOffset, ToPoint as _, + Anchor, Buffer, BufferSnapshot, ToOffset, ToPoint as _, language_settings::all_language_settings, }; use settings::EditPredictionPromptFormat; use std::{path::Path, sync::Arc, time::Instant}; -use zeta_prompt::ZetaPromptInput; +use zeta_prompt::{ZetaPromptInput, compute_editable_and_context_ranges}; const FIM_CONTEXT_TOKENS: usize = 512; @@ -62,34 +62,42 @@ pub fn request_prediction( let api_key = load_open_ai_compatible_api_key_if_needed(provider, cx); let result = cx.background_spawn(async move { - let (excerpt_range, _) = cursor_excerpt::editable_and_context_ranges_for_cursor_position( - cursor_point, - &snapshot, + let cursor_offset = cursor_point.to_offset(&snapshot); + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset); + let cursor_excerpt: Arc = snapshot + .text_for_range(excerpt_point_range.clone()) + .collect::() + .into(); + let syntax_ranges = + cursor_excerpt::compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range); + let (editable_range, _) = compute_editable_and_context_ranges( + &cursor_excerpt, + cursor_offset_in_excerpt, + &syntax_ranges, FIM_CONTEXT_TOKENS, 0, ); - let excerpt_offset_range = excerpt_range.to_offset(&snapshot); - let cursor_offset = cursor_point.to_offset(&snapshot); let inputs = ZetaPromptInput { events, related_files: Some(Vec::new()), cursor_offset_in_excerpt: cursor_offset - excerpt_offset_range.start, cursor_path: full_path.clone(), - excerpt_start_row: Some(excerpt_range.start.row), - cursor_excerpt: snapshot - .text_for_range(excerpt_range) - .collect::() - .into(), + excerpt_start_row: Some(excerpt_point_range.start.row), + cursor_excerpt, excerpt_ranges: Default::default(), + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, repo_url: None, }; - let prefix = inputs.cursor_excerpt[..inputs.cursor_offset_in_excerpt].to_string(); - let suffix = inputs.cursor_excerpt[inputs.cursor_offset_in_excerpt..].to_string(); + let editable_text = &inputs.cursor_excerpt[editable_range.clone()]; + let cursor_in_editable = cursor_offset_in_excerpt.saturating_sub(editable_range.start); + let prefix = editable_text[..cursor_in_editable].to_string(); + let suffix = editable_text[cursor_in_editable..].to_string(); let prompt = format_fim_prompt(prompt_format, &prefix, &suffix); let stop_tokens = get_fim_stop_tokens(); diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs index 0d63005feb18acb9a434ff107811080a7bcf1f12..1a88ba1f1f83a11f89ac282db46f91a3ee752f58 100644 --- a/crates/edit_prediction/src/mercury.rs +++ b/crates/edit_prediction/src/mercury.rs @@ -10,17 +10,14 @@ use gpui::{ App, AppContext as _, Entity, Global, SharedString, Task, http_client::{self, AsyncBody, HttpClient, Method}, }; -use language::{OffsetRangeExt as _, ToOffset, ToPoint as _}; +use language::{ToOffset, ToPoint as _}; use language_model::{ApiKeyState, EnvVar, env_var}; use release_channel::AppVersion; use serde::Serialize; use std::{mem, ops::Range, path::Path, sync::Arc, time::Instant}; - -use zeta_prompt::{ExcerptRanges, ZetaPromptInput}; +use zeta_prompt::ZetaPromptInput; const MERCURY_API_URL: &str = "https://api.inceptionlabs.ai/v1/edit/completions"; -const MAX_REWRITE_TOKENS: usize = 150; -const MAX_CONTEXT_TOKENS: usize = 350; pub struct Mercury { pub api_token: Entity, @@ -64,52 +61,46 @@ impl Mercury { let active_buffer = buffer.clone(); let result = cx.background_spawn(async move { - let (editable_range, context_range) = - crate::cursor_excerpt::editable_and_context_ranges_for_cursor_position( - cursor_point, - &snapshot, - MAX_CONTEXT_TOKENS, - MAX_REWRITE_TOKENS, - ); + let cursor_offset = cursor_point.to_offset(&snapshot); + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + crate::cursor_excerpt::compute_cursor_excerpt(&snapshot, cursor_offset); let related_files = zeta_prompt::filter_redundant_excerpts( related_files, full_path.as_ref(), - context_range.start.row..context_range.end.row, + excerpt_point_range.start.row..excerpt_point_range.end.row, ); - let context_offset_range = context_range.to_offset(&snapshot); - let context_start_row = context_range.start.row; - - let editable_offset_range = editable_range.to_offset(&snapshot); + let cursor_excerpt: Arc = snapshot + .text_for_range(excerpt_point_range.clone()) + .collect::() + .into(); + let syntax_ranges = crate::cursor_excerpt::compute_syntax_ranges( + &snapshot, + cursor_offset, + &excerpt_offset_range, + ); + let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges( + &cursor_excerpt, + cursor_offset_in_excerpt, + &syntax_ranges, + ); - let editable_range_in_excerpt = (editable_offset_range.start - - context_offset_range.start) - ..(editable_offset_range.end - context_offset_range.start); - let context_range_in_excerpt = - 0..(context_offset_range.end - context_offset_range.start); + let editable_offset_range = (excerpt_offset_range.start + + excerpt_ranges.editable_350.start) + ..(excerpt_offset_range.start + excerpt_ranges.editable_350.end); let inputs = zeta_prompt::ZetaPromptInput { events, related_files: Some(related_files), cursor_offset_in_excerpt: cursor_point.to_offset(&snapshot) - - context_offset_range.start, + - excerpt_offset_range.start, cursor_path: full_path.clone(), - cursor_excerpt: snapshot - .text_for_range(context_range) - .collect::() - .into(), + cursor_excerpt, experiment: None, - excerpt_start_row: Some(context_start_row), - excerpt_ranges: ExcerptRanges { - editable_150: editable_range_in_excerpt.clone(), - editable_180: editable_range_in_excerpt.clone(), - editable_350: editable_range_in_excerpt.clone(), - editable_150_context_350: context_range_in_excerpt.clone(), - editable_180_context_350: context_range_in_excerpt.clone(), - editable_350_context_150: context_range_in_excerpt.clone(), - ..Default::default() - }, + excerpt_start_row: Some(excerpt_point_range.start.row), + excerpt_ranges, + syntax_ranges: Some(syntax_ranges), in_open_source_repo: false, can_collect_data: false, repo_url: None, diff --git a/crates/edit_prediction/src/prediction.rs b/crates/edit_prediction/src/prediction.rs index 1c281453b93d0ab7c601f575b290c46fe63d2eae..ec4694c862be3ff1937dadc08ebab11b115b4cac 100644 --- a/crates/edit_prediction/src/prediction.rs +++ b/crates/edit_prediction/src/prediction.rs @@ -162,6 +162,7 @@ mod tests { cursor_excerpt: "".into(), excerpt_start_row: None, excerpt_ranges: Default::default(), + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, diff --git a/crates/edit_prediction/src/sweep_ai.rs b/crates/edit_prediction/src/sweep_ai.rs index ff5128e56e49191f308a574d5502f8139db9bc3f..d4e59885c86f44ceff98487940f2aaa435085e4d 100644 --- a/crates/edit_prediction/src/sweep_ai.rs +++ b/crates/edit_prediction/src/sweep_ai.rs @@ -226,6 +226,7 @@ impl SweepAi { editable_350_context_150: 0..inputs.snapshot.len(), ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 1a4d0b445a8c3d5876eb48646a0a1622a8b725a2..9362425c24df4199e40f97ac1278753c954a2f36 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -1,7 +1,8 @@ use crate::{ CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, StoredEvent, - ZedUpdateRequiredError, cursor_excerpt::compute_excerpt_ranges, + ZedUpdateRequiredError, + cursor_excerpt::{compute_cursor_excerpt, compute_syntax_ranges}, prediction::EditPredictionResult, }; use anyhow::Result; @@ -11,8 +12,7 @@ use cloud_llm_client::{ use edit_prediction_types::PredictedCursorPosition; use gpui::{App, AppContext as _, Entity, Task, WeakEntity, prelude::*}; use language::{ - Buffer, BufferSnapshot, ToOffset as _, ToPoint, language_settings::all_language_settings, - text_diff, + Buffer, BufferSnapshot, ToOffset as _, language_settings::all_language_settings, text_diff, }; use release_channel::AppVersion; use settings::EditPredictionPromptFormat; @@ -490,33 +490,35 @@ pub fn zeta2_prompt_input( can_collect_data: bool, repo_url: Option, ) -> (Range, zeta_prompt::ZetaPromptInput) { - let cursor_point = cursor_offset.to_point(snapshot); - - let (full_context, full_context_offset_range, excerpt_ranges) = - compute_excerpt_ranges(cursor_point, snapshot); - - let full_context_start_offset = full_context_offset_range.start; - let full_context_start_row = full_context.start.row; - - let cursor_offset_in_excerpt = cursor_offset - full_context_start_offset; + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + compute_cursor_excerpt(snapshot, cursor_offset); + + let cursor_excerpt: Arc = snapshot + .text_for_range(excerpt_point_range.clone()) + .collect::() + .into(); + let syntax_ranges = compute_syntax_ranges(snapshot, cursor_offset, &excerpt_offset_range); + let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges( + &cursor_excerpt, + cursor_offset_in_excerpt, + &syntax_ranges, + ); let prompt_input = zeta_prompt::ZetaPromptInput { cursor_path: excerpt_path, - cursor_excerpt: snapshot - .text_for_range(full_context) - .collect::() - .into(), + cursor_excerpt, cursor_offset_in_excerpt, - excerpt_start_row: Some(full_context_start_row), + excerpt_start_row: Some(excerpt_point_range.start.row), events, related_files: Some(related_files), excerpt_ranges, + syntax_ranges: Some(syntax_ranges), experiment: preferred_experiment, in_open_source_repo: is_open_source, can_collect_data, repo_url, }; - (full_context_offset_range, prompt_input) + (excerpt_offset_range, prompt_input) } pub(crate) fn edit_prediction_accepted( diff --git a/crates/edit_prediction_cli/src/load_project.rs b/crates/edit_prediction_cli/src/load_project.rs index f7e27ca432baacd38c468e5b4c6f97b62cb8ee3e..a9303451e8b6c6ae798be976a11d3b71fae99758 100644 --- a/crates/edit_prediction_cli/src/load_project.rs +++ b/crates/edit_prediction_cli/src/load_project.rs @@ -7,12 +7,12 @@ use crate::{ use anyhow::{Context as _, Result}; use edit_prediction::{ EditPredictionStore, - cursor_excerpt::compute_excerpt_ranges, + cursor_excerpt::{compute_cursor_excerpt, compute_syntax_ranges}, udiff::{OpenedBuffers, refresh_worktree_entries, strip_diff_path_prefix}, }; use futures::AsyncWriteExt as _; use gpui::{AsyncApp, Entity}; -use language::{Anchor, Buffer, LanguageNotFound, ToOffset, ToPoint}; +use language::{Anchor, Buffer, LanguageNotFound, ToOffset}; use project::{Project, ProjectPath, buffer_store::BufferStoreEvent}; use std::{fs, path::PathBuf, sync::Arc}; use zeta_prompt::ZetaPromptInput; @@ -75,32 +75,36 @@ pub async fn run_load_project( let (prompt_inputs, language_name) = buffer.read_with(&cx, |buffer, _cx| { let snapshot = buffer.snapshot(); - let cursor_point = cursor_position.to_point(&snapshot); let cursor_offset = cursor_position.to_offset(&snapshot); let language_name = buffer .language() .map(|l| l.name().to_string()) .unwrap_or_else(|| "Unknown".to_string()); - let (full_context_point_range, full_context_offset_range, excerpt_ranges) = - compute_excerpt_ranges(cursor_point, &snapshot); + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + compute_cursor_excerpt(&snapshot, cursor_offset); let cursor_excerpt: Arc = buffer - .text_for_range(full_context_offset_range.clone()) + .text_for_range(excerpt_offset_range.clone()) .collect::() .into(); - let cursor_offset_in_excerpt = cursor_offset - full_context_offset_range.start; - let excerpt_start_row = Some(full_context_point_range.start.row); + let syntax_ranges = compute_syntax_ranges(&snapshot, cursor_offset, &excerpt_offset_range); + let excerpt_ranges = zeta_prompt::compute_legacy_excerpt_ranges( + &cursor_excerpt, + cursor_offset_in_excerpt, + &syntax_ranges, + ); ( ZetaPromptInput { cursor_path: example.spec.cursor_path.clone(), cursor_excerpt, cursor_offset_in_excerpt, - excerpt_start_row, + excerpt_start_row: Some(excerpt_point_range.start.row), events, related_files: existing_related_files, excerpt_ranges, + syntax_ranges: Some(syntax_ranges), in_open_source_repo: false, can_collect_data: false, experiment: None, diff --git a/crates/edit_prediction_cli/src/retrieve_context.rs b/crates/edit_prediction_cli/src/retrieve_context.rs index 971bdf24d3e8cd1d8184a9009903cec25d3000d1..f02509ceb061db078d2a9a98b4322cf246b87594 100644 --- a/crates/edit_prediction_cli/src/retrieve_context.rs +++ b/crates/edit_prediction_cli/src/retrieve_context.rs @@ -24,6 +24,7 @@ pub async fn run_context_retrieval( .prompt_inputs .as_ref() .is_some_and(|inputs| inputs.related_files.is_some()) + || example.spec.repository_url.is_empty() { return Ok(()); } diff --git a/crates/edit_prediction_cli/src/reversal_tracking.rs b/crates/edit_prediction_cli/src/reversal_tracking.rs index 398ae24309bbb9368bb7947c94ad4f481c03ab9e..7623041a091acd3c726ba0281f090496255f8014 100644 --- a/crates/edit_prediction_cli/src/reversal_tracking.rs +++ b/crates/edit_prediction_cli/src/reversal_tracking.rs @@ -678,6 +678,7 @@ mod tests { editable_350_context_150: 0..content.len(), ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, diff --git a/crates/zeta_prompt/src/excerpt_ranges.rs b/crates/zeta_prompt/src/excerpt_ranges.rs new file mode 100644 index 0000000000000000000000000000000000000000..40621fe98a13bfa9195293ad29ba549240532a2e --- /dev/null +++ b/crates/zeta_prompt/src/excerpt_ranges.rs @@ -0,0 +1,443 @@ +use std::ops::Range; + +use serde::{Deserialize, Serialize}; + +use crate::estimate_tokens; + +/// Pre-computed byte offset ranges within `cursor_excerpt` for different +/// editable and context token budgets. Allows the server to select the +/// appropriate ranges for whichever model it uses. +#[derive(Clone, Debug, Default, PartialEq, Hash, Serialize, Deserialize)] +pub struct ExcerptRanges { + /// Editable region computed with a 150-token budget. + pub editable_150: Range, + /// Editable region computed with a 180-token budget. + pub editable_180: Range, + /// Editable region computed with a 350-token budget. + pub editable_350: Range, + /// Editable region computed with a 350-token budget. + pub editable_512: Option>, + /// Context boundary when using editable_150 with 350 tokens of additional context. + pub editable_150_context_350: Range, + /// Context boundary when using editable_180 with 350 tokens of additional context. + pub editable_180_context_350: Range, + /// Context boundary when using editable_350 with 150 tokens of additional context. + pub editable_350_context_150: Range, + pub editable_350_context_512: Option>, + pub editable_350_context_1024: Option>, + pub context_4096: Option>, + pub context_8192: Option>, +} + +/// Builds an `ExcerptRanges` by computing editable and context ranges for each +/// budget combination, using the syntax-aware logic in +/// `compute_editable_and_context_ranges`. +pub fn compute_legacy_excerpt_ranges( + cursor_excerpt: &str, + cursor_offset: usize, + syntax_ranges: &[Range], +) -> ExcerptRanges { + let compute = |editable_tokens, context_tokens| { + compute_editable_and_context_ranges( + cursor_excerpt, + cursor_offset, + syntax_ranges, + editable_tokens, + context_tokens, + ) + }; + + let (editable_150, editable_150_context_350) = compute(150, 350); + let (editable_180, editable_180_context_350) = compute(180, 350); + let (editable_350, editable_350_context_150) = compute(350, 150); + let (editable_512, _) = compute(512, 0); + let (_, editable_350_context_512) = compute(350, 512); + let (_, editable_350_context_1024) = compute(350, 1024); + let (_, context_4096) = compute(350, 4096); + let (_, context_8192) = compute(350, 8192); + + ExcerptRanges { + editable_150, + editable_180, + editable_350, + editable_512: Some(editable_512), + editable_150_context_350, + editable_180_context_350, + editable_350_context_150, + editable_350_context_512: Some(editable_350_context_512), + editable_350_context_1024: Some(editable_350_context_1024), + context_4096: Some(context_4096), + context_8192: Some(context_8192), + } +} + +/// Given the cursor excerpt text, cursor offset, and the syntax node ranges +/// containing the cursor (innermost to outermost), compute the editable range +/// and context range as byte offset ranges within `cursor_excerpt`. +/// +/// This is the server-side equivalent of `compute_excerpt_ranges` in +/// `edit_prediction::cursor_excerpt`, but operates on plain text with +/// pre-computed syntax boundaries instead of a `BufferSnapshot`. +pub fn compute_editable_and_context_ranges( + cursor_excerpt: &str, + cursor_offset: usize, + syntax_ranges: &[Range], + editable_token_limit: usize, + context_token_limit: usize, +) -> (Range, Range) { + let line_starts = compute_line_starts(cursor_excerpt); + let cursor_row = offset_to_row(&line_starts, cursor_offset); + let max_row = line_starts.len().saturating_sub(1) as u32; + + let editable_range = compute_editable_range_from_text( + cursor_excerpt, + &line_starts, + cursor_row, + max_row, + syntax_ranges, + editable_token_limit, + ); + + let context_range = expand_context_from_text( + cursor_excerpt, + &line_starts, + max_row, + &editable_range, + syntax_ranges, + context_token_limit, + ); + + (editable_range, context_range) +} + +fn compute_line_starts(text: &str) -> Vec { + let mut starts = vec![0]; + for (index, byte) in text.bytes().enumerate() { + if byte == b'\n' { + starts.push(index + 1); + } + } + starts +} + +fn offset_to_row(line_starts: &[usize], offset: usize) -> u32 { + match line_starts.binary_search(&offset) { + Ok(row) => row as u32, + Err(row) => (row.saturating_sub(1)) as u32, + } +} + +fn row_start_offset(line_starts: &[usize], row: u32) -> usize { + line_starts.get(row as usize).copied().unwrap_or(0) +} + +fn row_end_offset(text: &str, line_starts: &[usize], row: u32) -> usize { + if let Some(&next_start) = line_starts.get(row as usize + 1) { + // End before the newline of this row. + next_start.saturating_sub(1).min(text.len()) + } else { + text.len() + } +} + +fn row_range_to_byte_range( + text: &str, + line_starts: &[usize], + start_row: u32, + end_row: u32, +) -> Range { + let start = row_start_offset(line_starts, start_row); + let end = row_end_offset(text, line_starts, end_row); + start..end +} + +fn estimate_tokens_for_row_range( + text: &str, + line_starts: &[usize], + start_row: u32, + end_row: u32, +) -> usize { + let mut tokens = 0; + for row in start_row..end_row { + let row_len = row_end_offset(text, line_starts, row) + .saturating_sub(row_start_offset(line_starts, row)); + tokens += estimate_tokens(row_len).max(1); + } + tokens +} + +fn line_token_count_from_text(text: &str, line_starts: &[usize], row: u32) -> usize { + let row_len = + row_end_offset(text, line_starts, row).saturating_sub(row_start_offset(line_starts, row)); + estimate_tokens(row_len).max(1) +} + +/// Returns syntax boundaries (as row ranges) that contain the given row range +/// and extend beyond it, ordered from smallest to largest. +fn containing_syntax_boundaries_from_ranges( + line_starts: &[usize], + syntax_ranges: &[Range], + start_row: u32, + end_row: u32, +) -> Vec<(u32, u32)> { + let mut boundaries = Vec::new(); + let mut last: Option<(u32, u32)> = None; + + // syntax_ranges is innermost to outermost, so iterate in order. + for range in syntax_ranges { + let node_start_row = offset_to_row(line_starts, range.start); + let node_end_row = offset_to_row(line_starts, range.end); + + // Skip nodes that don't extend beyond the current range. + if node_start_row >= start_row && node_end_row <= end_row { + continue; + } + + let rows = (node_start_row, node_end_row); + if last == Some(rows) { + continue; + } + + last = Some(rows); + boundaries.push(rows); + } + + boundaries +} + +fn compute_editable_range_from_text( + text: &str, + line_starts: &[usize], + cursor_row: u32, + max_row: u32, + syntax_ranges: &[Range], + token_limit: usize, +) -> Range { + // Phase 1: Expand symmetrically from cursor using 75% of budget. + let initial_budget = (token_limit * 3) / 4; + let (mut start_row, mut end_row, mut remaining_tokens) = + expand_symmetric(text, line_starts, cursor_row, max_row, initial_budget); + + remaining_tokens += token_limit.saturating_sub(initial_budget); + + let original_start = start_row; + let original_end = end_row; + + // Phase 2: Expand to syntax boundaries that fit within budget. + let boundaries = + containing_syntax_boundaries_from_ranges(line_starts, syntax_ranges, start_row, end_row); + for (boundary_start, boundary_end) in &boundaries { + let tokens_for_start = if *boundary_start < start_row { + estimate_tokens_for_row_range(text, line_starts, *boundary_start, start_row) + } else { + 0 + }; + let tokens_for_end = if *boundary_end > end_row { + estimate_tokens_for_row_range(text, line_starts, end_row + 1, *boundary_end + 1) + } else { + 0 + }; + + let total_needed = tokens_for_start + tokens_for_end; + if total_needed <= remaining_tokens { + if *boundary_start < start_row { + start_row = *boundary_start; + } + if *boundary_end > end_row { + end_row = *boundary_end; + } + remaining_tokens = remaining_tokens.saturating_sub(total_needed); + } else { + break; + } + } + + // Phase 3: Continue line-wise in the direction we expanded least. + let expanded_up = original_start.saturating_sub(start_row); + let expanded_down = end_row.saturating_sub(original_end); + let prefer_up = expanded_up <= expanded_down; + + (start_row, end_row, _) = expand_linewise( + text, + line_starts, + start_row, + end_row, + max_row, + remaining_tokens, + prefer_up, + ); + + row_range_to_byte_range(text, line_starts, start_row, end_row) +} + +fn expand_context_from_text( + text: &str, + line_starts: &[usize], + max_row: u32, + editable_range: &Range, + syntax_ranges: &[Range], + context_token_limit: usize, +) -> Range { + let mut start_row = offset_to_row(line_starts, editable_range.start); + let mut end_row = offset_to_row(line_starts, editable_range.end); + let mut remaining_tokens = context_token_limit; + let mut did_syntax_expand = false; + + let boundaries = + containing_syntax_boundaries_from_ranges(line_starts, syntax_ranges, start_row, end_row); + for (boundary_start, boundary_end) in &boundaries { + let tokens_for_start = if *boundary_start < start_row { + estimate_tokens_for_row_range(text, line_starts, *boundary_start, start_row) + } else { + 0 + }; + let tokens_for_end = if *boundary_end > end_row { + estimate_tokens_for_row_range(text, line_starts, end_row + 1, *boundary_end + 1) + } else { + 0 + }; + + let total_needed = tokens_for_start + tokens_for_end; + if total_needed <= remaining_tokens { + if *boundary_start < start_row { + start_row = *boundary_start; + } + if *boundary_end > end_row { + end_row = *boundary_end; + } + remaining_tokens = remaining_tokens.saturating_sub(total_needed); + did_syntax_expand = true; + } else { + break; + } + } + + // Only expand line-wise if no syntax expansion occurred. + if !did_syntax_expand { + (start_row, end_row, _) = expand_linewise( + text, + line_starts, + start_row, + end_row, + max_row, + remaining_tokens, + true, + ); + } + + row_range_to_byte_range(text, line_starts, start_row, end_row) +} + +fn expand_symmetric( + text: &str, + line_starts: &[usize], + cursor_row: u32, + max_row: u32, + mut token_budget: usize, +) -> (u32, u32, usize) { + let mut start_row = cursor_row; + let mut end_row = cursor_row; + + let cursor_line_tokens = line_token_count_from_text(text, line_starts, cursor_row); + token_budget = token_budget.saturating_sub(cursor_line_tokens); + + loop { + let can_expand_up = start_row > 0; + let can_expand_down = end_row < max_row; + + if token_budget == 0 || (!can_expand_up && !can_expand_down) { + break; + } + + if can_expand_down { + let next_row = end_row + 1; + let line_tokens = line_token_count_from_text(text, line_starts, next_row); + if line_tokens <= token_budget { + end_row = next_row; + token_budget = token_budget.saturating_sub(line_tokens); + } else { + break; + } + } + + if can_expand_up && token_budget > 0 { + let next_row = start_row - 1; + let line_tokens = line_token_count_from_text(text, line_starts, next_row); + if line_tokens <= token_budget { + start_row = next_row; + token_budget = token_budget.saturating_sub(line_tokens); + } else { + break; + } + } + } + + (start_row, end_row, token_budget) +} + +fn expand_linewise( + text: &str, + line_starts: &[usize], + mut start_row: u32, + mut end_row: u32, + max_row: u32, + mut remaining_tokens: usize, + prefer_up: bool, +) -> (u32, u32, usize) { + loop { + let can_expand_up = start_row > 0; + let can_expand_down = end_row < max_row; + + if remaining_tokens == 0 || (!can_expand_up && !can_expand_down) { + break; + } + + let mut expanded = false; + + if prefer_up { + if can_expand_up { + let next_row = start_row - 1; + let line_tokens = line_token_count_from_text(text, line_starts, next_row); + if line_tokens <= remaining_tokens { + start_row = next_row; + remaining_tokens = remaining_tokens.saturating_sub(line_tokens); + expanded = true; + } + } + if can_expand_down && remaining_tokens > 0 { + let next_row = end_row + 1; + let line_tokens = line_token_count_from_text(text, line_starts, next_row); + if line_tokens <= remaining_tokens { + end_row = next_row; + remaining_tokens = remaining_tokens.saturating_sub(line_tokens); + expanded = true; + } + } + } else { + if can_expand_down { + let next_row = end_row + 1; + let line_tokens = line_token_count_from_text(text, line_starts, next_row); + if line_tokens <= remaining_tokens { + end_row = next_row; + remaining_tokens = remaining_tokens.saturating_sub(line_tokens); + expanded = true; + } + } + if can_expand_up && remaining_tokens > 0 { + let next_row = start_row - 1; + let line_tokens = line_token_count_from_text(text, line_starts, next_row); + if line_tokens <= remaining_tokens { + start_row = next_row; + remaining_tokens = remaining_tokens.saturating_sub(line_tokens); + expanded = true; + } + } + } + + if !expanded { + break; + } + } + + (start_row, end_row, remaining_tokens) +} diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 774ac7cb9baebb943c9223645aae8d16cd730998..87218f037f8ce61630b3a7505056d87f7af33376 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -1,3 +1,5 @@ +pub mod excerpt_ranges; + use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::fmt::Write; @@ -6,6 +8,10 @@ use std::path::Path; use std::sync::Arc; use strum::{EnumIter, IntoEnumIterator as _, IntoStaticStr}; +pub use crate::excerpt_ranges::{ + ExcerptRanges, compute_editable_and_context_ranges, compute_legacy_excerpt_ranges, +}; + pub const CURSOR_MARKER: &str = "<|user_cursor|>"; pub const MAX_PROMPT_TOKENS: usize = 4096; @@ -18,31 +24,6 @@ fn estimate_tokens(bytes: usize) -> usize { bytes / 3 } -/// Pre-computed byte offset ranges within `cursor_excerpt` for different -/// editable and context token budgets. Allows the server to select the -/// appropriate ranges for whichever model it uses. -#[derive(Clone, Debug, Default, PartialEq, Hash, Serialize, Deserialize)] -pub struct ExcerptRanges { - /// Editable region computed with a 150-token budget. - pub editable_150: Range, - /// Editable region computed with a 180-token budget. - pub editable_180: Range, - /// Editable region computed with a 350-token budget. - pub editable_350: Range, - /// Editable region computed with a 350-token budget. - pub editable_512: Option>, - /// Context boundary when using editable_150 with 350 tokens of additional context. - pub editable_150_context_350: Range, - /// Context boundary when using editable_180 with 350 tokens of additional context. - pub editable_180_context_350: Range, - /// Context boundary when using editable_350 with 150 tokens of additional context. - pub editable_350_context_150: Range, - pub editable_350_context_512: Option>, - pub editable_350_context_1024: Option>, - pub context_4096: Option>, - pub context_8192: Option>, -} - #[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)] pub struct ZetaPromptInput { pub cursor_path: Arc, @@ -55,6 +36,12 @@ pub struct ZetaPromptInput { pub related_files: Option>, /// These ranges let the server select model-appropriate subsets. pub excerpt_ranges: ExcerptRanges, + /// Byte offset ranges within `cursor_excerpt` for all syntax nodes that + /// contain `cursor_offset_in_excerpt`, ordered from innermost to outermost. + /// When present, the server uses these to compute editable/context ranges + /// instead of `excerpt_ranges`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub syntax_ranges: Option>>, /// The name of the edit prediction model experiment to use. #[serde(default, skip_serializing_if = "Option::is_none")] pub experiment: Option, @@ -223,6 +210,21 @@ pub fn special_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] } } +/// Returns the (editable_token_limit, context_token_limit) for a given format. +pub fn token_limits_for_format(format: ZetaFormat) -> (usize, usize) { + match format { + ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered => (150, 350), + ZetaFormat::V0114180EditableRegion => (180, 350), + ZetaFormat::V0120GitMergeMarkers + | ZetaFormat::V0131GitMergeMarkersPrefix + | ZetaFormat::V0211Prefill + | ZetaFormat::V0211SeedCoder + | ZetaFormat::v0226Hashline + | ZetaFormat::V0304SeedNoEdits => (350, 150), + ZetaFormat::V0304VariableEdit => (1024, 0), + } +} + pub fn stop_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] { match format { ZetaFormat::v0226Hashline => &[hashline::NO_EDITS_COMMAND_MARKER], @@ -262,8 +264,9 @@ pub fn excerpt_ranges_for_format( ), ZetaFormat::V0304VariableEdit => { let context = ranges - .context_8192 + .editable_350_context_1024 .clone() + .or(ranges.editable_350_context_512.clone()) .unwrap_or_else(|| ranges.editable_350_context_150.clone()); (context.clone(), context) } @@ -552,7 +555,18 @@ pub fn resolve_cursor_region( input: &ZetaPromptInput, format: ZetaFormat, ) -> (&str, Range, Range, usize) { - let (editable_range, context_range) = excerpt_range_for_format(format, &input.excerpt_ranges); + let (editable_range, context_range) = if let Some(syntax_ranges) = &input.syntax_ranges { + let (editable_tokens, context_tokens) = token_limits_for_format(format); + compute_editable_and_context_ranges( + &input.cursor_excerpt, + input.cursor_offset_in_excerpt, + syntax_ranges, + editable_tokens, + context_tokens, + ) + } else { + excerpt_range_for_format(format, &input.excerpt_ranges) + }; let context_start = context_range.start; let context_text = &input.cursor_excerpt[context_range.clone()]; let adjusted_editable = @@ -3876,6 +3890,7 @@ mod tests { editable_350_context_150: context_range, ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, @@ -3905,6 +3920,7 @@ mod tests { editable_350_context_150: context_range, ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, @@ -4488,6 +4504,7 @@ mod tests { editable_350_context_150: 0..excerpt.len(), ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, @@ -4551,6 +4568,7 @@ mod tests { editable_350_context_150: 0..28, ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false, @@ -4609,6 +4627,7 @@ mod tests { editable_350_context_150: context_range.clone(), ..Default::default() }, + syntax_ranges: None, experiment: None, in_open_source_repo: false, can_collect_data: false,