Detailed changes
@@ -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(
@@ -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()
@@ -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<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,
@@ -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::<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(
@@ -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,
@@ -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 =