diff --git a/crates/cloud_llm_client/src/cloud_llm_client.rs b/crates/cloud_llm_client/src/cloud_llm_client.rs index 212f8c548a0..152b84679da 100644 --- a/crates/cloud_llm_client/src/cloud_llm_client.rs +++ b/crates/cloud_llm_client/src/cloud_llm_client.rs @@ -1,5 +1,7 @@ #[cfg(feature = "predict-edits")] pub mod predict_edits_v3; +#[cfg(feature = "predict-edits")] +pub mod predict_edits_v4; use std::str::FromStr; use std::sync::Arc; diff --git a/crates/cloud_llm_client/src/predict_edits_v3.rs b/crates/cloud_llm_client/src/predict_edits_v3.rs index 77cdab33fee..2ffb32840d6 100644 --- a/crates/cloud_llm_client/src/predict_edits_v3.rs +++ b/crates/cloud_llm_client/src/predict_edits_v3.rs @@ -31,7 +31,7 @@ pub struct RawCompletionRequest { #[derive(Debug, Serialize, Deserialize)] pub struct PredictEditsV3Request { #[serde(flatten)] - pub input: zeta_prompt::ZetaPromptInput, + pub input: zeta_prompt::Zeta2PromptInput, } #[derive(Debug, Deserialize, Serialize)] diff --git a/crates/cloud_llm_client/src/predict_edits_v4.rs b/crates/cloud_llm_client/src/predict_edits_v4.rs new file mode 100644 index 00000000000..62f58f00bfe --- /dev/null +++ b/crates/cloud_llm_client/src/predict_edits_v4.rs @@ -0,0 +1,15 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct PredictEditsV4Request { + #[serde(flatten)] + pub input: zeta_prompt::Zeta3PromptInput, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct PredictEditsV4Response { + pub request_id: String, + pub patch: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_version: Option, +} diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 5e4cdf6da79..f0520bdf15f 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -59,7 +59,7 @@ use std::rc::Rc; use text::{AnchorRangeExt, Edit}; use workspace::{AppState, Workspace}; use zeta_prompt::ContextSource; -use zeta_prompt::{ZetaFormat, ZetaPromptInput}; +use zeta_prompt::{Zeta2PromptInput, ZetaFormat}; use std::mem; use std::ops::Range; @@ -99,9 +99,8 @@ use crate::license_detection::LicenseDetectionWatcher; use crate::mercury::Mercury; pub use crate::metrics::{KeptRateResult, compute_kept_rate}; use crate::onboarding_modal::ZedPredictModal; -pub use crate::prediction::EditPrediction; -pub use crate::prediction::EditPredictionId; use crate::prediction::EditPredictionResult; +pub use crate::prediction::{EditPrediction, EditPredictionId, EditPredictionInputs}; pub use language_model::ApiKeyState; pub use telemetry_events::EditPredictionRating; pub use zed_edit_prediction_delegate::ZedEditPredictionDelegate; @@ -2898,7 +2897,7 @@ impl EditPredictionStore { } pub(crate) async fn send_v3_request( - input: ZetaPromptInput, + input: Zeta2PromptInput, preferred_experiment: Option, client: Arc, llm_token: LlmApiToken, diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 7e4ec8abef4..9a6c47ffa2c 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -41,7 +41,7 @@ use util::{ }; use uuid::Uuid; use workspace::{AppState, CollaboratorId, MultiWorkspace}; -use zeta_prompt::ZetaPromptInput; +use zeta_prompt::Zeta2PromptInput; use crate::udiff::apply_diff_to_string; use crate::{ @@ -2699,7 +2699,7 @@ async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) { buffer: buffer.clone(), snapshot: cx.read(|cx| buffer.read(cx).snapshot()), id: EditPredictionId("the-id".into()), - inputs: ZetaPromptInput { + inputs: EditPredictionInputs::V2(Zeta2PromptInput { events: Default::default(), related_files: Default::default(), active_buffer_diagnostics: vec![], @@ -2712,7 +2712,7 @@ async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) { in_open_source_repo: false, can_collect_data: false, repo_url: None, - }, + }), model_version: None, trigger: PredictEditsRequestTrigger::Other, }; diff --git a/crates/edit_prediction/src/example_spec.rs b/crates/edit_prediction/src/example_spec.rs index c610b777dca..a0bb6bd2ff5 100644 --- a/crates/edit_prediction/src/example_spec.rs +++ b/crates/edit_prediction/src/example_spec.rs @@ -4,11 +4,10 @@ use std::{borrow::Cow, fmt::Write as _, mem, path::Path, sync::Arc}; use telemetry_events::EditPredictionRating; pub use zeta_prompt::udiff::{ - CURSOR_POSITION_MARKER, encode_cursor_in_patch, extract_cursor_from_patch, + CURSOR_POSITION_MARKER, INLINE_CURSOR_MARKER, encode_cursor_in_patch, extract_cursor_from_patch, }; use crate::data_collection::format_cursor_excerpt; -pub const INLINE_CURSOR_MARKER: &str = "<|user_cursor|>"; /// Maximum cursor file size to capture (64KB). /// Files larger than this will not have their content captured, @@ -494,8 +493,7 @@ impl ExampleSpec { /// to the start of the hunk. /// /// In the serialized representation of this example, the cursor position is represented - /// using a comment line in the diff, beginning with `#`, and containing a `[CURSOR_POSITION]` - /// marker with the same format as the [`Self::cursor_excerpt`]. + /// using an inline `<|user_cursor|>` marker in an added diff line. pub fn expected_patches_with_cursor_positions(&self) -> Vec<(String, Option)> { self.expected_patches .iter() @@ -784,8 +782,7 @@ mod tests { +// prints a greeting fn main() { - println!("hi"); - + println!("hello, {}", ); - # ^[CURSOR_POSITION] + + println!("hello, {}", <|user_cursor|>); let x = 42; } "#} @@ -814,8 +811,7 @@ mod tests { +++ b/test.rs @@ -1,2 +1,2 @@ -fn old() {} - +fn new_name() {} - # ^[CURSOR_POSITION] + +fn new_<|user_cursor|>name() {} "#}; let cursor_offset = "fn new_name() {}".find("name").unwrap(); @@ -826,7 +822,7 @@ mod tests { assert_eq!( encoded_once .lines() - .filter(|line| line.contains(CURSOR_POSITION_MARKER)) + .filter(|line| line.contains(INLINE_CURSOR_MARKER)) .count(), 1 ); diff --git a/crates/edit_prediction/src/fim.rs b/crates/edit_prediction/src/fim.rs index a28480a1141..349fc308318 100644 --- a/crates/edit_prediction/src/fim.rs +++ b/crates/edit_prediction/src/fim.rs @@ -1,5 +1,5 @@ use crate::{ - EditPredictionId, EditPredictionModelInput, cursor_excerpt, + EditPredictionId, EditPredictionInputs, EditPredictionModelInput, cursor_excerpt, open_ai_compatible::{self, load_open_ai_compatible_api_key_if_needed}, prediction::EditPredictionResult, }; @@ -10,7 +10,7 @@ use language::{ language_settings::all_language_settings, }; use std::{path::Path, sync::Arc, time::Instant}; -use zeta_prompt::{ZetaPromptInput, compute_editable_and_context_ranges}; +use zeta_prompt::{Zeta2PromptInput, compute_editable_and_context_ranges}; const FIM_CONTEXT_TOKENS: usize = 512; @@ -19,7 +19,7 @@ struct FimRequestOutput { edits: Vec<(std::ops::Range, Arc)>, editable_range: std::ops::Range, snapshot: BufferSnapshot, - inputs: ZetaPromptInput, + inputs: Zeta2PromptInput, buffer: Entity, } @@ -78,7 +78,7 @@ pub fn request_prediction( 0, ); - let inputs = ZetaPromptInput { + let inputs = Zeta2PromptInput { events, related_files: Some(Vec::new()), active_buffer_diagnostics: Vec::new(), @@ -154,7 +154,7 @@ pub fn request_prediction( output.edits.into(), None, Some(output.editable_range), - output.inputs, + EditPredictionInputs::V2(output.inputs), None, trigger, cx.background_executor().now() - request_start, diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs index e835b6ece84..986d708b110 100644 --- a/crates/edit_prediction/src/mercury.rs +++ b/crates/edit_prediction/src/mercury.rs @@ -1,7 +1,7 @@ use crate::{ - DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput, - EditPredictionStartedDebugEvent, EditPredictionStore, open_ai_response::text_from_response, - prediction::EditPredictionResult, zeta::compute_edits, + DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionInputs, + EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, + open_ai_response::text_from_response, prediction::EditPredictionResult, zeta::compute_edits, }; use anyhow::{Context as _, Result}; use cloud_llm_client::EditPredictionRejectReason; @@ -16,7 +16,7 @@ use language_model::{ApiKeyState, EnvVar, env_var}; use release_channel::AppVersion; use serde::{Deserialize, Serialize}; use std::{mem, ops::Range, path::Path, sync::Arc}; -use zeta_prompt::ZetaPromptInput; +use zeta_prompt::Zeta2PromptInput; const MERCURY_API_URL: &str = "https://api.inceptionlabs.ai/v1/edit/completions"; @@ -103,7 +103,7 @@ impl Mercury { + excerpt_ranges.editable_350.start) ..(excerpt_offset_range.start + excerpt_ranges.editable_350.end); - let inputs = zeta_prompt::ZetaPromptInput { + let inputs = zeta_prompt::Zeta2PromptInput { events, related_files: Some(related_files), cursor_offset_in_excerpt: cursor_point.to_offset(&snapshot) @@ -255,7 +255,7 @@ impl Mercury { edits.into(), None, Some(editable_range), - inputs, + EditPredictionInputs::V2(inputs), None, trigger, cx.background_executor().now() - request_start, @@ -267,7 +267,7 @@ impl Mercury { } } -fn build_prompt(inputs: &ZetaPromptInput) -> String { +fn build_prompt(inputs: &Zeta2PromptInput) -> String { const RECENTLY_VIEWED_SNIPPETS_START: &str = "<|recently_viewed_code_snippets|>\n"; const RECENTLY_VIEWED_SNIPPETS_END: &str = "<|/recently_viewed_code_snippets|>\n"; const RECENTLY_VIEWED_SNIPPET_START: &str = "<|recently_viewed_code_snippet|>\n"; diff --git a/crates/edit_prediction/src/prediction.rs b/crates/edit_prediction/src/prediction.rs index b9d41c44ea7..ca6379fa6ee 100644 --- a/crates/edit_prediction/src/prediction.rs +++ b/crates/edit_prediction/src/prediction.rs @@ -4,7 +4,7 @@ use cloud_llm_client::{EditPredictionRejectReason, PredictEditsRequestTrigger}; use edit_prediction_types::{PredictedCursorPosition, interpolate_edits}; use gpui::{AsyncApp, Entity, SharedString}; use language::{Anchor, Buffer, BufferSnapshot, EditPreview, TextBufferSnapshot}; -use zeta_prompt::ZetaPromptInput; +use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput}; #[derive(Clone, Default, Debug, PartialEq, Eq, Hash)] pub struct EditPredictionId(pub SharedString); @@ -21,6 +21,13 @@ impl std::fmt::Display for EditPredictionId { } } +#[derive(Clone, serde::Serialize)] +#[serde(tag = "version", content = "input")] +pub enum EditPredictionInputs { + V2(Zeta2PromptInput), + V3(Zeta3PromptInput), +} + /// A prediction response that was returned from the provider, whether it was ultimately valid or not. pub struct EditPredictionResult { pub prediction: EditPrediction, @@ -36,7 +43,7 @@ impl EditPredictionResult { edits: Arc<[(Range, Arc)]>, cursor_position: Option, editable_range: Option>, - inputs: ZetaPromptInput, + inputs: EditPredictionInputs, model_version: Option, trigger: PredictEditsRequestTrigger, e2e_latency: std::time::Duration, @@ -99,8 +106,8 @@ pub struct EditPrediction { pub editable_range: Option>, pub snapshot: BufferSnapshot, pub edit_preview: EditPreview, + pub inputs: EditPredictionInputs, pub buffer: Entity, - pub inputs: zeta_prompt::ZetaPromptInput, pub model_version: Option, pub trigger: PredictEditsRequestTrigger, } @@ -134,7 +141,7 @@ mod tests { use super::*; use gpui::{App, Entity, TestAppContext, prelude::*}; use language::{Buffer, ToOffset as _}; - use zeta_prompt::ZetaPromptInput; + use zeta_prompt::Zeta2PromptInput; #[gpui::test] async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) { @@ -157,7 +164,7 @@ mod tests { edit_preview, model_version: None, trigger: PredictEditsRequestTrigger::Other, - inputs: ZetaPromptInput { + inputs: EditPredictionInputs::V2(Zeta2PromptInput { events: vec![], related_files: Some(vec![]), active_buffer_diagnostics: vec![], @@ -170,7 +177,7 @@ mod tests { in_open_source_repo: false, can_collect_data: false, repo_url: None, - }, + }), }; cx.update(|cx| { diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 321ab27c027..cff9bb4d858 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -1,8 +1,8 @@ use crate::{ CloudRequestTimeoutError, CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, - EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent, - EditPredictionStore, PromptHistoryBoundary, ZedUpdateRequiredError, - buffer_path_with_id_fallback, + EditPredictionId, EditPredictionInputs, EditPredictionModelInput, + EditPredictionStartedDebugEvent, EditPredictionStore, PromptHistoryBoundary, + ZedUpdateRequiredError, buffer_path_with_id_fallback, cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, data_collection::CapturedPredictionContext, prediction::EditPredictionResult, @@ -24,7 +24,7 @@ use workspace::workspace_error::{ErrorAction, ErrorSeverity, WorkspaceError}; use std::{ops::Range, path::Path, sync::Arc}; use zeta_prompt::{ - ParsedOutput, ZetaFormat, ZetaPromptInput, excerpt_ranges_for_format, format_zeta_prompt, + ParsedOutput, Zeta2PromptInput, ZetaFormat, excerpt_ranges_for_format, format_zeta_prompt, get_prefill, parse_zeta2_model_output, stop_tokens_for_format, zeta1::{self, EDITABLE_REGION_END_MARKER}, }; @@ -84,7 +84,7 @@ pub(crate) fn request_prediction_with_zeta( let app_version = AppVersion::global(cx); struct Prediction { - prompt_input: ZetaPromptInput, + prompt_input: Zeta2PromptInput, buffer: Entity, snapshot: BufferSnapshot, edits: Vec<(Range, Arc)>, @@ -396,7 +396,7 @@ pub(crate) fn request_prediction_with_zeta( edits.into(), cursor_position, Some(editable_anchor_range), - inputs, + EditPredictionInputs::V2(inputs), model_version, trigger, request_duration, @@ -602,7 +602,7 @@ pub fn zeta2_prompt_input( is_open_source: bool, can_collect_data: bool, repo_url: Option, -) -> (Range, zeta_prompt::ZetaPromptInput) { +) -> (Range, zeta_prompt::Zeta2PromptInput) { let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = compute_cursor_excerpt(snapshot, cursor_offset); @@ -624,7 +624,7 @@ pub fn zeta2_prompt_input( ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT, ); - let prompt_input = zeta_prompt::ZetaPromptInput { + let prompt_input = zeta_prompt::Zeta2PromptInput { cursor_path: excerpt_path, cursor_excerpt, cursor_offset_in_excerpt, diff --git a/crates/edit_prediction_cli/src/example.rs b/crates/edit_prediction_cli/src/example.rs index bce5b2bc2a5..f7e2f298bf9 100644 --- a/crates/edit_prediction_cli/src/example.rs +++ b/crates/edit_prediction_cli/src/example.rs @@ -16,7 +16,7 @@ use std::{ io::Read, path::{Path, PathBuf}, }; -use zeta_prompt::ZetaPromptInput; +use zeta_prompt::Zeta2PromptInput; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct Example { @@ -26,7 +26,7 @@ pub struct Example { /// The full content of the file where an edit is being predicted, and the /// actual cursor offset. #[serde(skip_serializing_if = "Option::is_none")] - pub prompt_inputs: Option, + pub prompt_inputs: Option, /// The input and expected output from the edit prediction model. #[serde(skip_serializing_if = "Option::is_none")] diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs index 23c9c7dde11..b7a752f96b2 100644 --- a/crates/edit_prediction_cli/src/format_prompt.rs +++ b/crates/edit_prediction_cli/src/format_prompt.rs @@ -10,14 +10,14 @@ use gpui::AsyncApp; use std::ops::Range; use std::sync::Arc; use zeta_prompt::{ - ZetaFormat, ZetaPromptInput, format_edit_history_within_budget, format_expected_output, + Zeta2PromptInput, ZetaFormat, format_edit_history_within_budget, format_expected_output, format_zeta_prompt, hashed_regions::{self, SnippetMarkers}, max_edit_event_count_for_format, resolve_cursor_region, }; fn resolved_excerpt_ranges_for_format( - input: &zeta_prompt::ZetaPromptInput, + input: &zeta_prompt::Zeta2PromptInput, format: ZetaFormat, ) -> (Range, Range) { let (_, editable_range_in_context, context_range, _) = resolve_cursor_region(input, format); @@ -457,7 +457,7 @@ impl TeacherJumpsPrompt { Ok((patch, actual_cursor)) } - fn format_edit_history(prompt_inputs: &ZetaPromptInput) -> String { + fn format_edit_history(prompt_inputs: &Zeta2PromptInput) -> String { format_edit_history_within_budget( &prompt_inputs.events, "", @@ -473,7 +473,7 @@ impl TeacherJumpsPrompt { /// is skipped: it renders in its own prompt section via /// `format_cursor_excerpt`, and including it here would duplicate it. fn format_context( - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, marker_table: &[SnippetMarkers], max_tokens: usize, cursor_file_ix: usize, @@ -601,7 +601,7 @@ impl TeacherJumpsPrompt { /// content appears in the prompt exactly once. fn format_cursor_excerpt( example: &Example, - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, marker_table: &[SnippetMarkers], cursor: &hashed_regions::RelatedFileCursor, ) -> Result { @@ -841,7 +841,7 @@ mod tests { human_feedback: Vec::new(), rating: None, }, - prompt_inputs: Some(zeta_prompt::ZetaPromptInput { + prompt_inputs: Some(zeta_prompt::Zeta2PromptInput { cursor_path: std::path::Path::new("src/main.rs").into(), cursor_excerpt: cursor_excerpt.into(), cursor_offset_in_excerpt: cursor_offset, @@ -1598,7 +1598,7 @@ mod tests { .map(|index| format!("line{index:02}\n")) .collect::(); let cursor_offset = excerpt.find("line40").expect("cursor line exists"); - let prompt_inputs = zeta_prompt::ZetaPromptInput { + let prompt_inputs = zeta_prompt::Zeta2PromptInput { cursor_path: std::path::Path::new("src/main.rs").into(), cursor_excerpt: excerpt.clone().into(), cursor_offset_in_excerpt: cursor_offset, diff --git a/crates/edit_prediction_cli/src/load_project.rs b/crates/edit_prediction_cli/src/load_project.rs index f370f013ff0..cdb69fe09e9 100644 --- a/crates/edit_prediction_cli/src/load_project.rs +++ b/crates/edit_prediction_cli/src/load_project.rs @@ -15,7 +15,7 @@ use gpui::{AsyncApp, Entity}; use language::{Anchor, Buffer, LanguageNotFound, ToOffset}; use project::{Project, ProjectPath, buffer_store::BufferStoreEvent}; use std::{fs, path::PathBuf, sync::Arc}; -use zeta_prompt::ZetaPromptInput; +use zeta_prompt::Zeta2PromptInput; pub async fn run_load_project( example: &mut Example, @@ -96,7 +96,7 @@ pub async fn run_load_project( ); ( - ZetaPromptInput { + Zeta2PromptInput { cursor_path: example.spec.cursor_path.clone(), cursor_excerpt, cursor_offset_in_excerpt, diff --git a/crates/edit_prediction_cli/src/pull_examples.rs b/crates/edit_prediction_cli/src/pull_examples.rs index 15171c6feb9..2d3c0652def 100644 --- a/crates/edit_prediction_cli/src/pull_examples.rs +++ b/crates/edit_prediction_cli/src/pull_examples.rs @@ -12,7 +12,7 @@ use std::sync::Arc; use std::time::Duration; use telemetry_events::EditPredictionRating; -use zeta_prompt::{ZetaFormat, ZetaPromptInput, excerpt_range_for_format}; +use zeta_prompt::{Zeta2PromptInput, ZetaFormat, excerpt_range_for_format}; use crate::PredictionProvider; use crate::example::{Example, ExamplePrediction, ExamplePrompt}; @@ -1130,7 +1130,7 @@ fn rated_examples_from_response<'a>( let request_id = get_string("request_id"); let inputs_json = get_json("inputs"); - let inputs: Option = match &inputs_json { + let inputs: Option = match &inputs_json { Some(v) => match serde_json::from_value(v.clone()) { Ok(parsed) => Some(parsed), Err(e) => { @@ -1188,7 +1188,7 @@ fn build_rated_example( request_id: Option, device_id: String, time: String, - input: ZetaPromptInput, + input: Zeta2PromptInput, output: String, settled_editable_region: Option, rating: String, @@ -1296,7 +1296,7 @@ fn requested_examples_from_response<'a>( let device_id = get_string("device_id"); let time = get_string("time"); let input_json = get_json("input"); - let input: Option = + let input: Option = input_json.clone().and_then(|v| serde_json::from_value(v).ok()); let zed_version = get_string("zed_version"); @@ -1376,7 +1376,7 @@ fn settled_examples_from_response<'a>( let time = get_string("time"); let input_raw = get_value("input"); let input_json = parse_json_value(input_raw.as_ref()); - let input: Option = input_json + let input: Option = input_json .as_ref() .and_then(|parsed| serde_json::from_value(parsed.clone()).ok()); let requested_output = get_string("requested_output"); @@ -1497,7 +1497,7 @@ fn captured_examples_from_response<'a>( let time = get_string("time"); let input_raw = get_value("input"); let input_json = parse_json_value(input_raw.as_ref()); - let input: Option = input_json + let input: Option = input_json .as_ref() .and_then(|parsed| serde_json::from_value(parsed.clone()).ok()); let example_raw = get_value("example"); @@ -1580,7 +1580,7 @@ fn build_settled_example( request_id: String, device_id: String, time: String, - input: ZetaPromptInput, + input: Zeta2PromptInput, requested_output: String, settled_editable_region: String, requested_format: ZetaFormat, @@ -1636,7 +1636,7 @@ fn build_captured_example( request_id: String, device_id: String, time: String, - input: ZetaPromptInput, + input: Zeta2PromptInput, mut example_spec: ExampleSpec, settled_editable_region: String, zed_version: Option, @@ -1721,7 +1721,7 @@ fn rejected_examples_from_response<'a>( let device_id = get_string("device_id"); let time = get_string("time"); let input_json = get_json("input"); - let input: Option = + let input: Option = input_json.clone().and_then(|v| serde_json::from_value(v).ok()); let prompt = get_string("prompt"); let output = get_string("output"); @@ -1768,7 +1768,7 @@ fn build_rejected_example( request_id: String, device_id: String, time: String, - input: ZetaPromptInput, + input: Zeta2PromptInput, prompt: Option, output: String, settled_editable_region: Option, @@ -1872,7 +1872,7 @@ fn accepted_examples_from_response<'a>( let device_id = get_string("device_id"); let time = get_string("time"); let input_json = get_json("input"); - let input: Option = + let input: Option = input_json.clone().and_then(|v| serde_json::from_value(v).ok()); let prompt = get_string("prompt"); let output = get_string("output"); @@ -1913,7 +1913,7 @@ fn build_accepted_example( request_id: String, device_id: String, time: String, - input: ZetaPromptInput, + input: Zeta2PromptInput, prompt: Option, output: String, settled_editable_region: Option, @@ -1970,7 +1970,7 @@ fn build_example_from_snowflake( request_id: String, device_id: String, time: String, - input: ZetaPromptInput, + input: Zeta2PromptInput, tags: Vec, rejection: Option, zed_version: Option, diff --git a/crates/edit_prediction_cli/src/score.rs b/crates/edit_prediction_cli/src/score.rs index 5a0f763d049..e911f030273 100644 --- a/crates/edit_prediction_cli/src/score.rs +++ b/crates/edit_prediction_cli/src/score.rs @@ -202,7 +202,7 @@ pub fn run_context_coverage_scoring( fn context_excerpts( _example: &Example, - prompt_inputs: &zeta_prompt::ZetaPromptInput, + prompt_inputs: &zeta_prompt::Zeta2PromptInput, retrieved_context_byte_limit: Option, context_source_filter: Option<&[ContextSource]>, ) -> Vec { @@ -836,7 +836,7 @@ mod tests { use edit_prediction::example_spec::ExampleSpec; use edit_prediction_metrics::PredictionScore; use std::path::Path; - use zeta_prompt::{ExcerptRanges, RelatedExcerpt, ZetaPromptInput}; + use zeta_prompt::{ExcerptRanges, RelatedExcerpt, Zeta2PromptInput}; #[test] fn summary_includes_limited_filtered_retrieved_context_bytes_once_per_example() { @@ -889,7 +889,7 @@ mod tests { human_feedback: Vec::new(), rating: None, }, - prompt_inputs: Some(ZetaPromptInput { + prompt_inputs: Some(Zeta2PromptInput { cursor_path: Path::new("project/src/main.rs").into(), cursor_excerpt: "".into(), cursor_offset_in_excerpt: 0, diff --git a/crates/edit_prediction_metrics/src/reversal.rs b/crates/edit_prediction_metrics/src/reversal.rs index 8b1a6e5e634..80fdc40d53f 100644 --- a/crates/edit_prediction_metrics/src/reversal.rs +++ b/crates/edit_prediction_metrics/src/reversal.rs @@ -797,10 +797,10 @@ mod tests { use super::*; use indoc::indoc; use zeta_prompt::udiff::{apply_diff_to_string, unified_diff_with_context}; - use zeta_prompt::{ExcerptRanges, ZetaPromptInput}; + use zeta_prompt::{ExcerptRanges, Zeta2PromptInput}; fn compute_prediction_reversal_ratio( - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, predicted_content: &str, cursor_path: &Path, ) -> f32 { @@ -817,8 +817,8 @@ mod tests { content: &str, events: Vec>, excerpt_start_row: Option, - ) -> ZetaPromptInput { - ZetaPromptInput { + ) -> Zeta2PromptInput { + Zeta2PromptInput { cursor_path: Arc::from(Path::new("src/test.rs")), cursor_excerpt: content.into(), cursor_offset_in_excerpt: 0, diff --git a/crates/edit_prediction_ui/src/rate_prediction_modal.rs b/crates/edit_prediction_ui/src/rate_prediction_modal.rs index be19ac5a744..0dc4c1aaf29 100644 --- a/crates/edit_prediction_ui/src/rate_prediction_modal.rs +++ b/crates/edit_prediction_ui/src/rate_prediction_modal.rs @@ -1,6 +1,8 @@ use buffer_diff::BufferDiff; use cloud_llm_client::PredictEditsRequestTrigger; -use edit_prediction::{EditPrediction, EditPredictionRating, EditPredictionStore}; +use edit_prediction::{ + EditPrediction, EditPredictionInputs, EditPredictionRating, EditPredictionStore, +}; use editor::{Editor, Inlay, MultiBuffer}; use feature_flags::{FeatureFlag, PresenceFlag, register_feature_flag}; use gpui::{ @@ -18,13 +20,14 @@ use project::{ }; use settings::Settings as _; use std::rc::Rc; -use std::{fmt::Write, ops::Range, sync::Arc}; +use std::{fmt::Write, ops::Range, path::Path, sync::Arc}; use theme_settings::ThemeSettings; use ui::{ ContextMenu, DropdownMenu, KeyBinding, List, ListItem, ListItemSpacing, PopoverMenuHandle, Tooltip, prelude::*, }; use workspace::{ModalView, Workspace}; +use zeta_prompt::{ContextSource, FilePosition, RelatedExcerpt, RelatedFile, Zeta3PromptInput}; actions!( zeta, @@ -409,6 +412,145 @@ impl RatePredictionsModal { Some(format!("{header}{diff_body}")) } + fn write_formatted_inputs(formatted_inputs: &mut String, inputs: &EditPredictionInputs) { + match inputs { + EditPredictionInputs::V2(inputs) => { + Self::write_events(formatted_inputs, &inputs.events); + Self::write_related_files( + formatted_inputs, + inputs.related_files.as_deref().unwrap_or_default(), + ); + Self::write_cursor_excerpt( + formatted_inputs, + inputs.cursor_path.as_ref(), + inputs.cursor_excerpt.as_ref(), + inputs.cursor_offset_in_excerpt, + ); + } + EditPredictionInputs::V3(inputs) => { + Self::write_events(formatted_inputs, &inputs.events); + Self::write_related_files(formatted_inputs, &inputs.editable_context); + Self::write_zeta3_cursor_excerpt(formatted_inputs, inputs); + } + } + } + + fn write_events(formatted_inputs: &mut String, events: &[Arc]) { + write!(formatted_inputs, "## Events\n\n").unwrap(); + + for event in events { + formatted_inputs.push_str("```diff\n"); + zeta_prompt::write_event(formatted_inputs, event.as_ref()); + formatted_inputs.push_str("```\n\n"); + } + } + + fn write_related_files(formatted_inputs: &mut String, included_files: &[RelatedFile]) { + write!(formatted_inputs, "## Related files\n\n").unwrap(); + + for included_file in included_files { + write!(formatted_inputs, "### {}\n\n", included_file.path.display()).unwrap(); + + for excerpt in included_file.excerpts.iter() { + write!( + formatted_inputs, + "```{}\n{}\n```\n", + included_file.path.display(), + excerpt.text + ) + .unwrap(); + } + } + } + + fn write_zeta3_cursor_excerpt(formatted_inputs: &mut String, inputs: &Zeta3PromptInput) { + let current_excerpt = inputs + .editable_context + .iter() + .filter(|file| file.path == inputs.cursor_path) + .flat_map(|file| file.excerpts.iter()) + .find_map(|excerpt| { + if excerpt.context_source != ContextSource::CurrentFile { + return None; + } + + Some(( + excerpt, + Self::offset_for_position_in_excerpt(excerpt, inputs.cursor_position)?, + )) + }); + + if let Some((excerpt, cursor_offset)) = current_excerpt { + Self::write_cursor_excerpt( + formatted_inputs, + inputs.cursor_path.as_ref(), + excerpt.text.as_ref(), + cursor_offset, + ); + } else { + write!(formatted_inputs, "## Cursor Excerpt\n\n").unwrap(); + writeln!( + formatted_inputs, + "No current-file excerpt found for `{}` at row {}, column {}.", + inputs.cursor_path.display(), + inputs.cursor_position.row, + inputs.cursor_position.column + ) + .unwrap(); + } + } + + fn write_cursor_excerpt( + formatted_inputs: &mut String, + cursor_path: &Path, + cursor_excerpt: &str, + cursor_offset: usize, + ) { + write!(formatted_inputs, "## Cursor Excerpt\n\n").unwrap(); + + let mut cursor_offset = cursor_offset.min(cursor_excerpt.len()); + while !cursor_excerpt.is_char_boundary(cursor_offset) { + cursor_offset = cursor_offset.saturating_sub(1); + } + writeln!( + formatted_inputs, + "```{}\n{}{}\n```\n", + cursor_path.display(), + &cursor_excerpt[..cursor_offset], + &cursor_excerpt[cursor_offset..], + ) + .unwrap(); + } + + fn offset_for_position_in_excerpt( + excerpt: &RelatedExcerpt, + position: FilePosition, + ) -> Option { + if position.row < excerpt.row_range.start { + return None; + } + + let relative_row = (position.row - excerpt.row_range.start) as usize; + let text = excerpt.text.as_ref(); + let mut row_start = 0; + + for row in 0..=relative_row { + if row == relative_row { + let row_end = text[row_start..] + .find('\n') + .map_or(text.len(), |offset| row_start + offset); + let row_text = &text[row_start..row_end]; + let column = + row_text.floor_char_boundary((position.column as usize).min(row_text.len())); + return Some(row_start + column); + } + + row_start += text[row_start..].find('\n')? + 1; + } + + None + } + pub fn select_completion( &mut self, prediction: Option, @@ -524,63 +666,7 @@ impl RatePredictionsModal { }); let mut formatted_inputs = String::new(); - - write!(&mut formatted_inputs, "## Events\n\n").unwrap(); - - for event in &prediction.inputs.events { - formatted_inputs.push_str("```diff\n"); - zeta_prompt::write_event(&mut formatted_inputs, event.as_ref()); - formatted_inputs.push_str("```\n\n"); - } - - write!(&mut formatted_inputs, "## Related files\n\n").unwrap(); - - for included_file in prediction - .inputs - .related_files - .as_deref() - .unwrap_or_default() - .iter() - { - write!( - &mut formatted_inputs, - "### {}\n\n", - included_file.path.display() - ) - .unwrap(); - - for excerpt in included_file.excerpts.iter() { - write!( - &mut formatted_inputs, - "```{}\n{}\n```\n", - included_file.path.display(), - excerpt.text - ) - .unwrap(); - } - } - - write!(&mut formatted_inputs, "## Cursor Excerpt\n\n").unwrap(); - - let mut cursor_offset = prediction - .inputs - .cursor_offset_in_excerpt - .min(prediction.inputs.cursor_excerpt.len()); - while !prediction - .inputs - .cursor_excerpt - .is_char_boundary(cursor_offset) - { - cursor_offset = cursor_offset.saturating_sub(1); - } - writeln!( - &mut formatted_inputs, - "```{}\n{}{}\n```\n", - prediction.inputs.cursor_path.display(), - &prediction.inputs.cursor_excerpt[..cursor_offset], - &prediction.inputs.cursor_excerpt[cursor_offset..], - ) - .unwrap(); + Self::write_formatted_inputs(&mut formatted_inputs, &prediction.inputs); let current_editable_region = editable_range.as_ref().map(|range| { prediction diff --git a/crates/zeta_prompt/src/hashed_regions.rs b/crates/zeta_prompt/src/hashed_regions.rs index 59b927e46fa..44c9ca8bab6 100644 --- a/crates/zeta_prompt/src/hashed_regions.rs +++ b/crates/zeta_prompt/src/hashed_regions.rs @@ -13,7 +13,7 @@ //! weren't run through current-file retrieval can be normalized with //! [`ensure_cursor_file_excerpt`] before rendering or parsing. -use crate::{ContextSource, RelatedExcerpt, RelatedFile, ZetaPromptInput, multi_region, udiff}; +use crate::{ContextSource, RelatedExcerpt, RelatedFile, Zeta2PromptInput, multi_region, udiff}; use anyhow::{Context as _, Result, anyhow}; use std::{ borrow::Cow, @@ -50,11 +50,11 @@ pub struct SnippetMarkers { /// /// The assignment is deterministic and independent of any later budget-based /// truncation, so the same table can be rebuilt when parsing model output. -pub fn build_marker_table(input: &ZetaPromptInput) -> Vec { +pub fn build_marker_table(input: &Zeta2PromptInput) -> Vec { build_marker_table_with_filter(input, |_| true) } -pub fn build_editable_marker_table(input: &ZetaPromptInput) -> Vec { +pub fn build_editable_marker_table(input: &Zeta2PromptInput) -> Vec { build_marker_table_with_filter(input, is_hash_region_editable_context_source) } @@ -66,7 +66,7 @@ pub fn is_hash_region_editable_context_source(context_source: ContextSource) -> } fn build_marker_table_with_filter( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, include_context_source: impl Fn(ContextSource) -> bool, ) -> Vec { let mut used_ids = HashSet::new(); @@ -274,7 +274,7 @@ fn line_start_offset(text: &str, row: usize) -> Option { Some(offset) } -pub fn locate_cursor_in_related_files(input: &ZetaPromptInput) -> Option { +pub fn locate_cursor_in_related_files(input: &Zeta2PromptInput) -> Option { let related_files = input.related_files.as_deref()?; let excerpt_start_row = input.excerpt_start_row?; let cursor_offset = input.cursor_offset_in_excerpt; @@ -316,7 +316,7 @@ pub fn locate_cursor_in_related_files(input: &ZetaPromptInput) -> Option Option bool { +pub fn ensure_cursor_file_excerpt(input: &mut Zeta2PromptInput) -> bool { if locate_cursor_in_related_files(input).is_some() { return true; } @@ -391,7 +391,7 @@ pub fn marker_table_for_excerpt( } fn merge_contiguous_snippets( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, marker_table: Vec, ) -> Result>> { let related_files = input @@ -441,7 +441,7 @@ fn merge_contiguous_snippets( } fn snippet_path_and_start_row( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, snippet: &ParseSnippet<'_>, ) -> Result<(PathBuf, u32)> { let related_files = input @@ -565,7 +565,7 @@ pub fn encode_from_old_and_new( /// which emits exactly two tags per block and no intermediate tags. Any /// unpaired trailing tag is ignored. pub fn parse_output_as_patch( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, output: &str, cursor_marker: &str, ) -> Result { @@ -651,7 +651,7 @@ fn find_all_marker_tags(text: &str) -> Vec<(String, usize, usize)> { /// every region that contains it; the returned [`HashRegionCursor`] reports the /// first such position. pub fn build_patch_from_spans( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, spans: &[(String, String, String)], cursor_marker: &str, ) -> Result<(String, Option)> { @@ -731,7 +731,7 @@ pub fn build_patch_from_spans( /// Apply resolved edits to their snippets and emit one diff section per edited /// snippet, in the order snippets first appear in the edit sequence. fn assemble_patch_from_edits( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, snippets: &[ParseSnippet<'_>], edits: Vec, ) -> Result<(String, Option)> { @@ -887,7 +887,7 @@ fn detect_trailing_deletion(old_span: &str, new_span: &str) -> Option { /// If at least one hunk is reachable, the remaining hunks are still encoded /// (partial edit). If no hunk is reachable, the output is `NO_EDITS`. pub fn encode_patch_as_output( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, patch: &str, cursor_offset: Option, cursor_marker: &str, @@ -973,8 +973,8 @@ mod tests { use std::path::PathBuf; use std::sync::Arc; - fn make_input(cursor_excerpt: &str, related: &[(&str, &[&str])]) -> ZetaPromptInput { - ZetaPromptInput { + fn make_input(cursor_excerpt: &str, related: &[(&str, &[&str])]) -> Zeta2PromptInput { + Zeta2PromptInput { cursor_path: PathBuf::from("src/main.rs").into(), cursor_excerpt: cursor_excerpt.into(), cursor_offset_in_excerpt: 0, diff --git a/crates/zeta_prompt/src/udiff.rs b/crates/zeta_prompt/src/udiff.rs index eee1e65931c..7b4f01e9adf 100644 --- a/crates/zeta_prompt/src/udiff.rs +++ b/crates/zeta_prompt/src/udiff.rs @@ -88,53 +88,37 @@ pub fn strip_diff_metadata(diff: &str) -> String { result } -/// Marker used to encode cursor position in patch comment lines. pub const CURSOR_POSITION_MARKER: &str = "[CURSOR_POSITION]"; +pub const INLINE_CURSOR_MARKER: &str = "<|user_cursor|>"; /// Extract cursor offset from a patch and return `(clean_patch, cursor_offset)`. -/// -/// Cursor position is encoded as a comment line (starting with `#`) containing -/// `[CURSOR_POSITION]`. A `^` in the line indicates the cursor column; a `<` -/// indicates column 0. The offset is computed relative to addition (`+`) and -/// context (` `) lines accumulated so far in the hunk, which represent the -/// cursor position within the new text contributed by the hunk. pub fn extract_cursor_from_patch(patch: &str) -> (String, Option) { let mut clean_patch = String::new(); - let mut cursor_offset: Option = None; + let mut cursor_offset = None; let mut line_start_offset = 0usize; - let mut prev_line_start_offset = 0usize; for line in patch.lines() { - let diff_line = DiffLine::parse(line); + if !clean_patch.is_empty() { + clean_patch.push('\n'); + } - match &diff_line { - DiffLine::Garbage(content) - if content.starts_with('#') && content.contains(CURSOR_POSITION_MARKER) => - { - let caret_column = if let Some(caret_pos) = content.find('^') { - caret_pos - } else if content.find('<').is_some() { - 0 - } else { - continue; - }; - let cursor_column = caret_column.saturating_sub('#'.len_utf8()); - cursor_offset = Some(prev_line_start_offset + cursor_column); - } - _ => { - if !clean_patch.is_empty() { - clean_patch.push('\n'); + match DiffLine::parse(line) { + DiffLine::Addition(content) => { + let clean_content = content.replace(INLINE_CURSOR_MARKER, ""); + if cursor_offset.is_none() + && let Some(marker_offset) = content.find(INLINE_CURSOR_MARKER) + { + cursor_offset = Some(line_start_offset + marker_offset); } + clean_patch.push('+'); + clean_patch.push_str(&clean_content); + line_start_offset += clean_content.len() + 1; + } + DiffLine::Context(content) => { clean_patch.push_str(line); - - match diff_line { - DiffLine::Addition(content) | DiffLine::Context(content) => { - prev_line_start_offset = line_start_offset; - line_start_offset += content.len() + 1; - } - _ => {} - } + line_start_offset += content.len() + 1; } + _ => clean_patch.push_str(line), } } @@ -356,40 +340,33 @@ pub fn encode_cursor_in_patch(patch: &str, cursor_offset: Option) -> Stri let mut line_start_offset = 0usize; for line in patch.lines() { - if matches!( - DiffLine::parse(line), - DiffLine::Garbage(content) - if content.starts_with('#') && content.contains(CURSOR_POSITION_MARKER) - ) { - continue; - } - if !result.is_empty() { result.push('\n'); } - result.push_str(line); match DiffLine::parse(line) { DiffLine::Addition(content) => { + let content = content.replace(INLINE_CURSOR_MARKER, ""); let line_end_offset = line_start_offset + content.len(); - - if cursor_offset >= line_start_offset && cursor_offset <= line_end_offset { - let cursor_column = cursor_offset - line_start_offset; - - result.push('\n'); - result.push('#'); - for _ in 0..cursor_column { - result.push(' '); - } - write!(result, "^{}", CURSOR_POSITION_MARKER).unwrap(); + result.push('+'); + if cursor_offset >= line_start_offset + && cursor_offset <= line_end_offset + && let Some(before) = content.get(..cursor_offset - line_start_offset) + && let Some(after) = content.get(cursor_offset - line_start_offset..) + { + result.push_str(before); + result.push_str(INLINE_CURSOR_MARKER); + result.push_str(after); + } else { + result.push_str(&content); } - line_start_offset = line_end_offset + 1; } DiffLine::Context(content) => { + result.push_str(line); line_start_offset += content.len() + 1; } - _ => {} + _ => result.push_str(line), } } diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 5f106d01b5e..446bd0ff689 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -51,7 +51,7 @@ pub fn clamp_text_to_token_count(text: &str, max_tokens: usize) -> &str { } #[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)] -pub struct ZetaPromptInput { +pub struct Zeta2PromptInput { pub cursor_path: Arc, pub cursor_excerpt: Arc, pub cursor_offset_in_excerpt: usize, @@ -78,6 +78,30 @@ pub struct ZetaPromptInput { pub repo_url: Option, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct FilePosition { + pub row: u32, + pub column: u32, +} + +#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)] +pub struct Zeta3PromptInput { + pub cursor_path: Arc, + pub cursor_position: FilePosition, + pub events: Vec>, + pub editable_context: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub syntax_ranges: Vec>, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub active_buffer_diagnostics: Vec, + #[serde(default)] + pub in_open_source_repo: bool, + #[serde(default)] + pub can_collect_data: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub repo_url: Option, +} + #[derive( Default, Clone, @@ -275,7 +299,7 @@ pub enum ContextSource { OracleSnippet, } -pub fn prompt_input_contains_special_tokens(input: &ZetaPromptInput, format: ZetaFormat) -> bool { +pub fn prompt_input_contains_special_tokens(input: &Zeta2PromptInput, format: ZetaFormat) -> bool { special_tokens_for_format(format).iter().any(|token| { if let Some(line_token) = token.strip_suffix('\n') { input.cursor_excerpt.lines().any(|line| line == line_token) @@ -285,8 +309,49 @@ pub fn prompt_input_contains_special_tokens(input: &ZetaPromptInput, format: Zet }) } -pub fn format_zeta_prompt(input: &ZetaPromptInput, format: ZetaFormat) -> Option { - let max_prompt_tokens = match format { +pub fn format_zeta_prompt(input: &Zeta2PromptInput, format: ZetaFormat) -> Option { + format_prompt_with_budget_for_format(input, format, max_prompt_tokens_for_format(format)) +} + +pub fn format_zeta3_prompt(input: &Zeta3PromptInput, format: ZetaFormat) -> Option { + match format { + ZetaFormat::V0318SeedMultiRegions => {} + _ => return None, + } + + let (current_excerpt, cursor_offset_in_excerpt) = zeta3_current_file_excerpt(input)?; + let (context, editable_range, context_range, cursor_offset) = resolve_zeta3_cursor_region( + current_excerpt.text.as_ref(), + cursor_offset_in_excerpt, + &input.syntax_ranges, + format, + ); + let relative_row_range = + offset_range_to_row_range(current_excerpt.text.as_ref(), context_range); + let cursor_row_range = current_excerpt.row_range.start + relative_row_range.start + ..current_excerpt.row_range.start + relative_row_range.end; + let related_files = filter_redundant_excerpts( + zeta3_related_files(input, current_excerpt), + input.cursor_path.as_ref(), + cursor_row_range, + ); + + format_resolved_prompt_with_budget( + format, + input.cursor_path.as_ref(), + context, + &editable_range, + cursor_offset, + &input.events, + &related_files, + &input.active_buffer_diagnostics, + Some(input.cursor_position.row), + max_prompt_tokens_for_format(format), + ) +} + +fn max_prompt_tokens_for_format(format: ZetaFormat) -> usize { + match format { ZetaFormat::V0112MiddleAtEnd | ZetaFormat::V0113Ordered | ZetaFormat::V0114180EditableRegion @@ -306,9 +371,88 @@ pub fn format_zeta_prompt(input: &ZetaPromptInput, format: ZetaFormat) -> Option ZetaFormat::V0615HashRegions => 8000, ZetaFormat::V0420Diagnostics => 8192, ZetaFormat::V0327SingleFile => 16384, - }; + } +} - format_prompt_with_budget_for_format(input, format, max_prompt_tokens) +fn zeta3_current_file_excerpt(input: &Zeta3PromptInput) -> Option<(&RelatedExcerpt, usize)> { + input + .editable_context + .iter() + .filter(|file| file.path == input.cursor_path) + .flat_map(|file| file.excerpts.iter()) + .find_map(|excerpt| { + if excerpt.context_source != ContextSource::CurrentFile { + return None; + } + Some(( + excerpt, + offset_for_position_in_excerpt(excerpt, input.cursor_position)?, + )) + }) +} + +fn offset_for_position_in_excerpt( + excerpt: &RelatedExcerpt, + position: FilePosition, +) -> Option { + if position.row < excerpt.row_range.start { + return None; + } + + let relative_row = (position.row - excerpt.row_range.start) as usize; + let text = excerpt.text.as_ref(); + let mut row_start = 0; + + for row in 0..=relative_row { + if row == relative_row { + let row_end = text[row_start..] + .find('\n') + .map_or(text.len(), |offset| row_start + offset); + let row_text = &text[row_start..row_end]; + let column = + row_text.floor_char_boundary((position.column as usize).min(row_text.len())); + return Some(row_start + column); + } + + row_start += text[row_start..].find('\n')? + 1; + } + + None +} + +fn zeta3_related_files( + input: &Zeta3PromptInput, + current_excerpt: &RelatedExcerpt, +) -> Vec { + input + .editable_context + .iter() + .filter_map(|file| { + let mut file = file.clone(); + if file.path == input.cursor_path { + file.excerpts.retain(|excerpt| excerpt != current_excerpt); + } + (!file.excerpts.is_empty()).then_some(file) + }) + .collect() +} + +fn resolve_zeta3_cursor_region<'a>( + cursor_excerpt: &'a str, + cursor_offset: usize, + syntax_ranges: &[Range], + format: ZetaFormat, +) -> (&'a str, Range, Range, usize) { + let (editable_tokens, context_tokens) = token_limits_for_format(format); + let (editable_range, context_range) = compute_editable_and_context_ranges( + cursor_excerpt, + cursor_offset, + syntax_ranges, + editable_tokens, + context_tokens, + ); + + adjust_cursor_region(cursor_excerpt, cursor_offset, editable_range, context_range) } pub fn special_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] { @@ -888,7 +1032,7 @@ fn assemble_single_file_fim_prompt( } fn format_hash_region_related_files_within_budget( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, marker_table: &[hashed_regions::SnippetMarkers], cursor: &hashed_regions::RelatedFileCursor, max_tokens: usize, @@ -1020,7 +1164,7 @@ fn format_hash_region_related_files_within_budget( } fn format_hash_regions_prompt_with_budget( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, max_tokens: usize, ) -> Option { let marker_table = hashed_regions::build_marker_table(input); @@ -1054,14 +1198,19 @@ fn format_hash_regions_prompt_with_budget( } pub fn format_prompt_with_budget_for_format( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, format: ZetaFormat, max_tokens: usize, ) -> Option { + if format == ZetaFormat::V0615HashRegions { + return format_hash_regions_prompt_with_budget( + input, + apply_prompt_budget_margin(max_tokens), + ); + } + let (context, editable_range, context_range, cursor_offset) = resolve_cursor_region(input, format); - let path = &*input.cursor_path; - let empty_files = Vec::new(); let input_related_files = input.related_files.as_deref().unwrap_or(&empty_files); let filtered_related_files = if format == ZetaFormat::V0615HashRegions { @@ -1079,12 +1228,41 @@ pub fn format_prompt_with_budget_for_format( } else { input_related_files.to_vec() }; - let related_files = filtered_related_files.as_slice(); + let cursor_buffer_row = input.excerpt_start_row.map(|excerpt_start_row| { + excerpt_start_row + + input.cursor_excerpt[..context_range.start + cursor_offset] + .bytes() + .filter(|byte| *byte == b'\n') + .count() as u32 + }); + format_resolved_prompt_with_budget( + format, + input.cursor_path.as_ref(), + context, + &editable_range, + cursor_offset, + &input.events, + &filtered_related_files, + &input.active_buffer_diagnostics, + cursor_buffer_row, + max_tokens, + ) +} + +fn format_resolved_prompt_with_budget( + format: ZetaFormat, + path: &Path, + context: &str, + editable_range: &Range, + cursor_offset: usize, + events: &[Arc], + related_files: &[RelatedFile], + active_buffer_diagnostics: &[ActiveBufferDiagnostic], + cursor_buffer_row: Option, + max_tokens: usize, +) -> Option { let prompt = match format { - ZetaFormat::V0615HashRegions => { - format_hash_regions_prompt_with_budget(input, apply_prompt_budget_margin(max_tokens))? - } ZetaFormat::V0211SeedCoder | ZetaFormat::V0331SeedCoderModelPy | ZetaFormat::V0304SeedNoEdits @@ -1100,27 +1278,19 @@ pub fn format_prompt_with_budget_for_format( &mut cursor_section, path, context, - &editable_range, + editable_range, cursor_offset, ); - let cursor_buffer_row = input.excerpt_start_row.map(|excerpt_start_row| { - excerpt_start_row - + input.cursor_excerpt[..context_range.start + cursor_offset] - .bytes() - .filter(|byte| *byte == b'\n') - .count() as u32 - }); - let budget_with_margin = apply_prompt_budget_margin(max_tokens); seed_coder::assemble_fim_prompt( context, - &editable_range, + editable_range, &cursor_section, - &input.events, + events, related_files, if format == ZetaFormat::V0420Diagnostics { - &input.active_buffer_diagnostics + active_buffer_diagnostics } else { &[] }, @@ -1136,15 +1306,15 @@ pub fn format_prompt_with_budget_for_format( &mut cursor_section, path, context, - &editable_range, + editable_range, cursor_offset, ); qwen::assemble_fim_prompt( context, - &editable_range, + editable_range, &cursor_section, - &input.events, + events, related_files, apply_prompt_budget_margin(max_tokens), ) @@ -1156,15 +1326,15 @@ pub fn format_prompt_with_budget_for_format( &mut cursor_section, path, context, - &editable_range, + editable_range, cursor_offset, ); assemble_single_file_fim_prompt( context, - &editable_range, + editable_range, &cursor_section, - &input.events, + events, apply_prompt_budget_margin(max_tokens), ) } @@ -1175,7 +1345,7 @@ pub fn format_prompt_with_budget_for_format( &mut cursor_section, path, context, - &editable_range, + editable_range, cursor_offset, ); @@ -1184,7 +1354,7 @@ pub fn format_prompt_with_budget_for_format( remaining_budget = remaining_budget.saturating_sub(cursor_tokens); let edit_history_section = format_edit_history_within_budget( - &input.events, + events, "<|file_sep|>", "edit history", remaining_budget, @@ -1194,7 +1364,7 @@ pub fn format_prompt_with_budget_for_format( remaining_budget = remaining_budget.saturating_sub(edit_history_tokens); let related_files_section = format_related_files_within_budget( - &related_files, + related_files, "<|file_sep|>", "", remaining_budget, @@ -1473,10 +1643,10 @@ pub fn encode_patch_as_output_for_format( } } -/// Given a `ZetaPromptInput`, a format, and a patch (with cursor already +/// Given a `Zeta2PromptInput`, a format, and a patch (with cursor already /// extracted), produce the expected model output string for training. pub fn format_expected_output( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, format: ZetaFormat, patch: &str, cursor_offset: Option, @@ -1659,7 +1829,7 @@ pub fn parsed_output_from_editable_region( pub fn parse_zeta2_model_output( output: &str, format: ZetaFormat, - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, ) -> Result { let output = match output_end_marker_for_format(format) { Some(marker) => output.strip_suffix(marker).unwrap_or(output), @@ -1737,7 +1907,7 @@ pub fn parse_zeta2_model_output( pub fn parse_zeta2_model_output_as_patch( output: &str, format: ZetaFormat, - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, ) -> Result { if format == ZetaFormat::V0615HashRegions { return hashed_regions::parse_output_as_patch(prompt_inputs, output, CURSOR_MARKER); @@ -1747,8 +1917,46 @@ pub fn parse_zeta2_model_output_as_patch( parsed_output_to_patch(prompt_inputs, parsed) } +pub fn parse_zeta3_model_output_as_patch( + output: &str, + format: ZetaFormat, + input: &Zeta3PromptInput, +) -> Result { + match format { + ZetaFormat::V0318SeedMultiRegions => {} + _ => anyhow::bail!("unsupported Zeta3 output format: {format}"), + } + + let output = match output_end_marker_for_format(format) { + Some(marker) => output.strip_suffix(marker).unwrap_or(output), + None => output, + }; + let (current_excerpt, cursor_offset_in_excerpt) = zeta3_current_file_excerpt(input) + .ok_or_else(|| anyhow!("Zeta3 input is missing current-file editable context at cursor"))?; + let (context, editable_range_in_context, context_range, _) = resolve_zeta3_cursor_region( + current_excerpt.text.as_ref(), + cursor_offset_in_excerpt, + &input.syntax_ranges, + format, + ); + let old_editable_region = &context[editable_range_in_context.clone()]; + let range_in_excerpt = editable_range_in_context.start + context_range.start + ..editable_range_in_context.end + context_range.start; + let parsed = parsed_output_from_editable_region( + range_in_excerpt, + multi_region::apply_marker_span_v0318(old_editable_region, output)?, + ); + + parsed_output_to_patch_for_excerpt( + input.cursor_path.as_ref(), + current_excerpt.text.as_ref(), + current_excerpt.row_range.start, + parsed, + ) +} + pub fn cursor_position_from_parsed_output( - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, parsed: &ParsedOutput, ) -> Option { let cursor_offset = parsed.cursor_offset_in_new_editable_region?; @@ -1785,11 +1993,24 @@ pub fn cursor_position_from_parsed_output( } pub fn parsed_output_to_patch( - prompt_inputs: &ZetaPromptInput, + prompt_inputs: &Zeta2PromptInput, + parsed: ParsedOutput, +) -> Result { + parsed_output_to_patch_for_excerpt( + prompt_inputs.cursor_path.as_ref(), + prompt_inputs.cursor_excerpt.as_ref(), + 0, + parsed, + ) +} + +fn parsed_output_to_patch_for_excerpt( + path: &Path, + excerpt: &str, + excerpt_start_row: u32, parsed: ParsedOutput, ) -> Result { let range_in_excerpt = parsed.range_in_excerpt; - let excerpt = prompt_inputs.cursor_excerpt.as_ref(); let old_text = excerpt[range_in_excerpt.clone()].to_string(); let mut new_text = parsed.new_editable_region; @@ -1802,7 +2023,8 @@ pub fn parsed_output_to_patch( } let editable_region_offset = range_in_excerpt.start; - let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count() as u32; + let editable_region_start_line = + excerpt_start_row + excerpt[..editable_region_offset].matches('\n').count() as u32; let editable_region_lines = old_text_normalized.lines().count() as u32; let diff = udiff::unified_diff_with_context( @@ -1813,11 +2035,7 @@ pub fn parsed_output_to_patch( editable_region_lines, ); - let path = prompt_inputs - .cursor_path - .to_string_lossy() - .trim_start_matches('/') - .to_string(); + let path = path.to_string_lossy().trim_start_matches('/').to_string(); let formatted_diff = format!("--- a/{path}\n+++ b/{path}\n{diff}"); Ok(udiff::encode_cursor_in_patch( @@ -1834,7 +2052,7 @@ pub fn excerpt_range_for_format( } pub fn resolve_cursor_region( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, format: ZetaFormat, ) -> (&str, Range, Range, usize) { let (editable_range, context_range) = if format == ZetaFormat::V0327SingleFile { @@ -1859,11 +2077,25 @@ pub fn resolve_cursor_region( excerpt_range_for_format(format, &input.excerpt_ranges) }; + adjust_cursor_region( + &input.cursor_excerpt, + input.cursor_offset_in_excerpt, + editable_range, + context_range, + ) +} + +fn adjust_cursor_region( + cursor_excerpt: &str, + cursor_offset: usize, + editable_range: Range, + context_range: Range, +) -> (&str, Range, Range, usize) { let context_start = context_range.start; - let context_text = &input.cursor_excerpt[context_range.clone()]; + let context_text = &cursor_excerpt[context_range.clone()]; let adjusted_editable = (editable_range.start - context_start)..(editable_range.end - context_start); - let adjusted_cursor = input.cursor_offset_in_excerpt - context_start; + let adjusted_cursor = cursor_offset - context_start; ( context_text, @@ -1873,7 +2105,7 @@ pub fn resolve_cursor_region( ) } -pub fn get_prefill(input: &ZetaPromptInput, format: ZetaFormat) -> String { +pub fn get_prefill(input: &Zeta2PromptInput, format: ZetaFormat) -> String { let (context, editable_range, _, _) = resolve_cursor_region(input, format); get_prefill_for_format(format, context, &editable_range) } @@ -5120,10 +5352,10 @@ pub mod zeta1 { prompt } - /// Formats a complete zeta1 prompt from a `ZetaPromptInput` using the given + /// Formats a complete zeta1 prompt from a `Zeta2PromptInput` using the given /// editable and context byte-offset ranges within `cursor_excerpt`. pub fn format_zeta1_from_input( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, editable_range: Range, context_range: Range, ) -> String { @@ -5189,7 +5421,7 @@ pub mod zeta1 { /// Formats the excerpt section of a zeta1 prompt using byte-offset ranges /// within `cursor_excerpt`. fn format_zeta1_excerpt( - input: &ZetaPromptInput, + input: &Zeta2PromptInput, editable_range: Range, context_range: Range, ) -> String { @@ -5298,9 +5530,9 @@ mod tests { cursor_offset: usize, events: Vec, related_files: Vec, - ) -> ZetaPromptInput { + ) -> Zeta2PromptInput { let context_range = 0..cursor_excerpt.len(); - ZetaPromptInput { + Zeta2PromptInput { cursor_path: Path::new("test.rs").into(), cursor_excerpt: cursor_excerpt.into(), cursor_offset_in_excerpt: cursor_offset, @@ -5329,8 +5561,8 @@ mod tests { editable_range: Range, context_range: Range, cursor_offset: usize, - ) -> ZetaPromptInput { - ZetaPromptInput { + ) -> Zeta2PromptInput { + Zeta2PromptInput { cursor_path: Path::new("test.rs").into(), cursor_excerpt: excerpt.into(), cursor_offset_in_excerpt: cursor_offset, @@ -5380,7 +5612,7 @@ mod tests { } } - fn format_with_budget(input: &ZetaPromptInput, max_tokens: usize) -> Option { + fn format_with_budget(input: &Zeta2PromptInput, max_tokens: usize) -> Option { format_prompt_with_budget_for_format(input, ZetaFormat::V0114180EditableRegion, max_tokens) } @@ -5811,13 +6043,13 @@ mod tests { } #[track_caller] - fn format_seed_coder(input: &ZetaPromptInput) -> String { + fn format_seed_coder(input: &Zeta2PromptInput) -> String { format_prompt_with_budget_for_format(input, ZetaFormat::V0211SeedCoder, 10000) .expect("seed coder prompt formatting should succeed") } #[track_caller] - fn format_seed_coder_with_budget(input: &ZetaPromptInput, max_tokens: usize) -> String { + fn format_seed_coder_with_budget(input: &Zeta2PromptInput, max_tokens: usize) -> String { format_prompt_with_budget_for_format(input, ZetaFormat::V0211SeedCoder, max_tokens) .expect("seed coder prompt formatting should succeed") } @@ -6056,7 +6288,7 @@ mod tests { fn test_v0615_formats_hashed_markers_for_rendered_related_context() { let current_text = "fn main() {\n let value = 1;\n}\n"; let cursor_offset = current_text.find("let value").unwrap(); - let input = ZetaPromptInput { + let input = Zeta2PromptInput { cursor_path: Path::new("test.rs").into(), cursor_excerpt: current_text.into(), cursor_offset_in_excerpt: cursor_offset, @@ -6130,7 +6362,7 @@ mod tests { fn test_v0615_parse_related_file_jump_as_patch() { let current_text = "fn main() {\n helper();\n}\n"; let helper_text = "fn helper() {\n one();\n}\n"; - let input = ZetaPromptInput { + let input = Zeta2PromptInput { cursor_path: Path::new("test.rs").into(), cursor_excerpt: current_text.into(), cursor_offset_in_excerpt: current_text.find("helper").unwrap(), @@ -6190,7 +6422,7 @@ mod tests { fn test_v0615_expected_output_round_trips_to_patch() { let current_text = "fn main() {\n helper();\n}\n"; let helper_text = "fn helper() {\n one();\n}\n"; - let input = ZetaPromptInput { + let input = Zeta2PromptInput { cursor_path: Path::new("test.rs").into(), cursor_excerpt: current_text.into(), cursor_offset_in_excerpt: current_text.find("helper").unwrap(), @@ -6254,6 +6486,110 @@ mod tests { ); } + #[test] + fn test_zeta3_prompt_matches_zeta2_seed_multi_region_prompt() { + let excerpt = "fn main() {\n helper();\n}\n"; + let cursor_offset = excerpt.find("helper").expect("cursor text exists") + "help".len(); + let cursor_row_start = excerpt.find(" helper").expect("cursor row exists"); + let syntax_ranges = vec![0..excerpt.len()]; + let related_file = make_related_file("related.rs", "fn helper() {}\n"); + let mut zeta2_input = make_input( + excerpt, + 0..excerpt.len(), + cursor_offset, + vec![make_event("test.rs", "-old\n+new\n")], + vec![related_file.clone()], + ); + zeta2_input.excerpt_start_row = Some(10); + zeta2_input.excerpt_ranges = + compute_legacy_excerpt_ranges(excerpt, cursor_offset, &syntax_ranges); + zeta2_input.syntax_ranges = Some(syntax_ranges.clone()); + + let zeta3_input = Zeta3PromptInput { + cursor_path: Path::new("test.rs").into(), + cursor_position: FilePosition { + row: 11, + column: (cursor_offset - cursor_row_start) as u32, + }, + events: vec![Arc::new(make_event("test.rs", "-old\n+new\n"))], + editable_context: vec![ + RelatedFile { + path: Path::new("test.rs").into(), + max_row: 12, + excerpts: vec![RelatedExcerpt { + row_range: 10..12, + text: excerpt.into(), + order: 0, + context_source: ContextSource::CurrentFile, + }], + in_open_source_repo: false, + }, + related_file, + ], + syntax_ranges, + active_buffer_diagnostics: vec![], + in_open_source_repo: false, + can_collect_data: false, + repo_url: None, + }; + + assert_eq!( + format_zeta3_prompt(&zeta3_input, ZetaFormat::V0318SeedMultiRegions), + format_zeta_prompt(&zeta2_input, ZetaFormat::V0318SeedMultiRegions), + ); + } + + #[test] + fn test_zeta3_parse_seed_multi_region_output_as_patch() { + let prefix = (0..20) + .map(|index| format!("prefix {index}\n")) + .collect::(); + let excerpt = "fn main() {\n let value = 1;\n dbg!(value);\n}\n"; + let new_excerpt = "fn main() {\n let value = 2;\n dbg!(value);\n}\n"; + let cursor_offset = excerpt.find("value =").expect("cursor text exists"); + let input = Zeta3PromptInput { + cursor_path: Path::new("test.rs").into(), + cursor_position: FilePosition { row: 21, column: 8 }, + events: vec![], + editable_context: vec![RelatedFile { + path: Path::new("test.rs").into(), + max_row: 24, + excerpts: vec![RelatedExcerpt { + row_range: 20..24, + text: excerpt.into(), + order: 0, + context_source: ContextSource::CurrentFile, + }], + in_open_source_repo: false, + }], + syntax_ranges: vec![0..excerpt.len()], + active_buffer_diagnostics: vec![], + in_open_source_repo: false, + can_collect_data: false, + repo_url: None, + }; + let output = multi_region::encode_from_old_and_new_v0318( + excerpt, + new_excerpt, + Some(cursor_offset), + CURSOR_MARKER, + multi_region::V0318_END_MARKER, + ) + .unwrap(); + + let patch = + parse_zeta3_model_output_as_patch(&output, ZetaFormat::V0318SeedMultiRegions, &input) + .unwrap(); + let full_text = format!("{prefix}{excerpt}"); + let expected = format!("{prefix}{new_excerpt}"); + + assert!(patch.contains(CURSOR_MARKER)); + assert_eq!( + udiff::apply_diff_to_string(&patch.replace(CURSOR_MARKER, ""), &full_text).unwrap(), + expected + ); + } + #[test] fn test_seed_coder_no_context() { let input = make_input("before\nmiddle\nafter", 7..13, 10, vec![], vec![]); @@ -6391,7 +6727,7 @@ mod tests { #[test] fn test_format_zeta1_from_input_basic() { let excerpt = "fn before() {}\nfn foo() {\n let x = 1;\n}\nfn after() {}\n"; - let input = ZetaPromptInput { + let input = Zeta2PromptInput { cursor_path: Path::new("src/main.rs").into(), cursor_excerpt: excerpt.into(), cursor_offset_in_excerpt: 30, @@ -6455,7 +6791,7 @@ mod tests { #[test] fn test_format_zeta1_from_input_no_start_of_file() { let excerpt = "fn foo() {\n let x = 1;\n}\n"; - let input = ZetaPromptInput { + let input = Zeta2PromptInput { cursor_path: Path::new("src/main.rs").into(), cursor_excerpt: excerpt.into(), cursor_offset_in_excerpt: 15, @@ -6514,7 +6850,7 @@ mod tests { let editable_range = 10..37; let context_range = 0..excerpt.len(); - let input = ZetaPromptInput { + let input = Zeta2PromptInput { cursor_path: Path::new("test.rs").into(), cursor_excerpt: excerpt.into(), cursor_offset_in_excerpt: 25,