From 831de8e48fd39e079f413ff2673dc072d7cd2c48 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Tue, 23 Sep 2025 16:50:07 -0300 Subject: [PATCH] zeta2: Include edits in prompt and add `max_prompt_bytes` param (#38737) Release Notes: - N/A Co-authored-by: Michael Sloan --- Cargo.lock | 2 + .../cloud_llm_client/src/predict_edits_v3.rs | 2 + crates/cloud_zeta2_prompt/Cargo.toml | 1 + .../src/cloud_zeta2_prompt.rs | 89 ++++++++++++++++--- crates/zeta2/Cargo.toml | 1 + crates/zeta2/src/zeta2.rs | 8 +- crates/zeta2_tools/src/zeta2_tools.rs | 35 +++++--- crates/zeta_cli/src/main.rs | 13 +-- 8 files changed, 118 insertions(+), 33 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b34c98fa72cc8cd2739fe24577af67a940cac10f..d2e287de62aff60ceeb7ccd313db603c303690d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3233,6 +3233,7 @@ version = "0.1.0" dependencies = [ "anyhow", "cloud_llm_client", + "indoc", "ordered-float 2.10.1", "rustc-hash 2.1.1", "strum 0.27.1", @@ -21651,6 +21652,7 @@ dependencies = [ "chrono", "client", "cloud_llm_client", + "cloud_zeta2_prompt", "edit_prediction", "edit_prediction_context", "futures 0.3.31", diff --git a/crates/cloud_llm_client/src/predict_edits_v3.rs b/crates/cloud_llm_client/src/predict_edits_v3.rs index 90f2a8b24fda6f38d52a21c90d184f9a66e1e80d..eeca7ed4e24d594d4a7e9555b3d9fbfd6f3706d2 100644 --- a/crates/cloud_llm_client/src/predict_edits_v3.rs +++ b/crates/cloud_llm_client/src/predict_edits_v3.rs @@ -29,8 +29,10 @@ pub struct PredictEditsRequest { /// Info about the git repository state, only present when can_collect_data is true. #[serde(skip_serializing_if = "Option::is_none", default)] pub git_info: Option, + // Only available to staff #[serde(default)] pub debug_info: bool, + pub prompt_max_bytes: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/cloud_zeta2_prompt/Cargo.toml b/crates/cloud_zeta2_prompt/Cargo.toml index a1194f13615964fd3013eb8dbdf3057984946e32..b06431baf652c67459280ffec64d45ca072420ef 100644 --- a/crates/cloud_zeta2_prompt/Cargo.toml +++ b/crates/cloud_zeta2_prompt/Cargo.toml @@ -14,6 +14,7 @@ path = "src/cloud_zeta2_prompt.rs" [dependencies] anyhow.workspace = true cloud_llm_client.workspace = true +indoc.workspace = true ordered-float.workspace = true rustc-hash.workspace = true strum.workspace = true diff --git a/crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs b/crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs index 6690380c74b0d4880210b683f34eea1d98a7946b..89cfb4c41f2d05c521c81c6e136d7bc46861d7bc 100644 --- a/crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs +++ b/crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs @@ -1,17 +1,28 @@ //! Zeta2 prompt planning and generation code shared with cloud. use anyhow::{Result, anyhow}; -use cloud_llm_client::predict_edits_v3::{self, ReferencedDeclaration}; +use cloud_llm_client::predict_edits_v3::{self, Event, ReferencedDeclaration}; +use indoc::indoc; use ordered_float::OrderedFloat; use rustc_hash::{FxHashMap, FxHashSet}; +use std::fmt::Write; use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path}; use strum::{EnumIter, IntoEnumIterator}; +pub const DEFAULT_MAX_PROMPT_BYTES: usize = 10 * 1024; + pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>"; /// NOTE: Differs from zed version of constant - includes a newline -pub const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>\n"; +pub const EDITABLE_REGION_START_MARKER_WITH_NEWLINE: &str = "<|editable_region_start|>\n"; /// NOTE: Differs from zed version of constant - includes a newline -pub const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>\n"; +pub const EDITABLE_REGION_END_MARKER_WITH_NEWLINE: &str = "<|editable_region_end|>\n"; + +// TODO: use constants for markers? +pub const SYSTEM_PROMPT: &str = indoc! {" + You are a code completion assistant and your task is to analyze user edits and then rewrite an excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking into account the cursor location. + + The excerpt to edit will be wrapped in markers <|editable_region_start|> and <|editable_region_end|>. The cursor position is marked with <|user_cursor_is_here|>. Please respond with edited code for that region. +"}; pub struct PlannedPrompt<'a> { request: &'a predict_edits_v3::PredictEditsRequest, @@ -286,7 +297,7 @@ impl<'a> PlannedPrompt<'a> { let mut excerpt_file_insertions = vec![ ( self.request.excerpt_range.start, - EDITABLE_REGION_START_MARKER, + EDITABLE_REGION_START_MARKER_WITH_NEWLINE, ), ( self.request.excerpt_range.start + self.request.cursor_offset, @@ -298,10 +309,67 @@ impl<'a> PlannedPrompt<'a> { .end .saturating_sub(0) .max(self.request.excerpt_range.start), - EDITABLE_REGION_END_MARKER, + EDITABLE_REGION_END_MARKER_WITH_NEWLINE, ), ]; + let mut output = String::new(); + output.push_str("## User Edits\n\n"); + Self::push_events(&mut output, &self.request.events); + + output.push_str("\n## Code\n\n"); + Self::push_file_snippets(&mut output, &mut excerpt_file_insertions, file_snippets); + output + } + + fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) { + for event in events { + match event { + Event::BufferChange { + path, + old_path, + diff, + predicted, + } => { + if let Some(old_path) = &old_path + && let Some(new_path) = &path + { + if old_path != new_path { + writeln!( + output, + "User renamed {} to {}\n\n", + old_path.display(), + new_path.display() + ) + .unwrap(); + } + } + + let path = path + .as_ref() + .map_or_else(|| "untitled".to_string(), |path| path.display().to_string()); + + if *predicted { + writeln!( + output, + "User accepted prediction {:?}:\n```diff\n{}\n```\n", + path, diff + ) + .unwrap(); + } else { + writeln!(output, "User edited {:?}:\n```diff\n{}\n```\n", path, diff) + .unwrap(); + } + } + } + } + } + + fn push_file_snippets( + output: &mut String, + excerpt_file_insertions: &mut Vec<(usize, &'static str)>, + file_snippets: Vec<(&Path, Vec<&PlannedSnippet>, bool)>, + ) { fn push_excerpt_file_range( range: Range, text: &str, @@ -325,7 +393,6 @@ impl<'a> PlannedPrompt<'a> { output.push_str(&text[last_offset - range.start..]); } - let mut output = String::new(); for (file_path, mut snippets, is_excerpt_file) in file_snippets { output.push_str(&format!("```{}\n", file_path.display())); @@ -345,8 +412,8 @@ impl<'a> PlannedPrompt<'a> { push_excerpt_file_range( last_range.end..snippet.range.end, text, - &mut excerpt_file_insertions, - &mut output, + excerpt_file_insertions, + output, ); } else { output.push_str(text); @@ -361,8 +428,8 @@ impl<'a> PlannedPrompt<'a> { push_excerpt_file_range( snippet.range.clone(), snippet.text, - &mut excerpt_file_insertions, - &mut output, + excerpt_file_insertions, + output, ); } else { output.push_str(snippet.text); @@ -372,8 +439,6 @@ impl<'a> PlannedPrompt<'a> { output.push_str("```\n\n"); } - - output } } diff --git a/crates/zeta2/Cargo.toml b/crates/zeta2/Cargo.toml index d362441cdd522e012d75bec9bde6daeff50e3e89..11ca5fcfddd3e7a7c1085401e7575bc599aea91e 100644 --- a/crates/zeta2/Cargo.toml +++ b/crates/zeta2/Cargo.toml @@ -17,6 +17,7 @@ arrayvec.workspace = true chrono.workspace = true client.workspace = true cloud_llm_client.workspace = true +cloud_zeta2_prompt.workspace = true edit_prediction.workspace = true edit_prediction_context.workspace = true futures.workspace = true diff --git a/crates/zeta2/src/zeta2.rs b/crates/zeta2/src/zeta2.rs index 4319f946f6cbf7319826e22594ed6331b2726889..47afc797d5d1d4dafca71f39ba15fbe6fd621c20 100644 --- a/crates/zeta2/src/zeta2.rs +++ b/crates/zeta2/src/zeta2.rs @@ -6,6 +6,7 @@ use cloud_llm_client::predict_edits_v3::{self, Signature}; use cloud_llm_client::{ EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME, ZED_VERSION_HEADER_NAME, }; +use cloud_zeta2_prompt::DEFAULT_MAX_PROMPT_BYTES; use edit_prediction::{DataCollectionState, Direction, EditPredictionProvider}; use edit_prediction_context::{ DeclarationId, EditPredictionContext, EditPredictionExcerptOptions, SyntaxIndex, @@ -49,6 +50,7 @@ pub const DEFAULT_EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPrediction pub const DEFAULT_OPTIONS: ZetaOptions = ZetaOptions { excerpt: DEFAULT_EXCERPT_OPTIONS, + max_prompt_bytes: DEFAULT_MAX_PROMPT_BYTES, max_diagnostic_bytes: 2048, }; @@ -71,6 +73,7 @@ pub struct Zeta { #[derive(Debug, Clone, PartialEq)] pub struct ZetaOptions { pub excerpt: EditPredictionExcerptOptions, + pub max_prompt_bytes: usize, pub max_diagnostic_bytes: usize, } @@ -408,6 +411,7 @@ impl Zeta { debug_context.is_some(), &worktree_snapshots, index_state.as_deref(), + Some(options.max_prompt_bytes), ); let retrieval_time = chrono::Utc::now() - before_retrieval; @@ -702,6 +706,7 @@ impl Zeta { debug_info, &worktree_snapshots, index_state.as_deref(), + Some(options.max_prompt_bytes), ) }) }) @@ -1062,6 +1067,7 @@ fn make_cloud_request( debug_info: bool, worktrees: &Vec, index_state: Option<&SyntaxIndexState>, + prompt_max_bytes: Option, ) -> predict_edits_v3::PredictEditsRequest { let mut signatures = Vec::new(); let mut declaration_to_signature_index = HashMap::default(); @@ -1132,9 +1138,9 @@ fn make_cloud_request( can_collect_data, diagnostic_groups, diagnostic_groups_truncated, - git_info, debug_info, + prompt_max_bytes, } } diff --git a/crates/zeta2_tools/src/zeta2_tools.rs b/crates/zeta2_tools/src/zeta2_tools.rs index de2cc660bef9f1326ee95eb145cd8b3c551a7acf..2dfa292c4380bd7295dbe50669237f0bf865cb7e 100644 --- a/crates/zeta2_tools/src/zeta2_tools.rs +++ b/crates/zeta2_tools/src/zeta2_tools.rs @@ -57,13 +57,16 @@ pub fn init(cx: &mut App) { .detach(); } +// TODO show included diagnostics, and events + pub struct Zeta2Inspector { focus_handle: FocusHandle, project: Entity, last_prediction: Option, - max_bytes_input: Entity, - min_bytes_input: Entity, + max_excerpt_bytes_input: Entity, + min_excerpt_bytes_input: Entity, cursor_context_ratio_input: Entity, + max_prompt_bytes_input: Entity, active_view: ActiveView, zeta: Entity, _active_editor_subscription: Option, @@ -129,9 +132,10 @@ impl Zeta2Inspector { project: project.clone(), last_prediction: None, active_view: ActiveView::Context, - max_bytes_input: Self::number_input("Max Bytes", window, cx), - min_bytes_input: Self::number_input("Min Bytes", window, cx), + max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx), + min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx), cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx), + max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx), zeta: zeta.clone(), _active_editor_subscription: None, _update_state_task: Task::ready(()), @@ -147,10 +151,10 @@ impl Zeta2Inspector { window: &mut Window, cx: &mut Context, ) { - self.max_bytes_input.update(cx, |input, cx| { + self.max_excerpt_bytes_input.update(cx, |input, cx| { input.set_text(options.excerpt.max_bytes.to_string(), window, cx); }); - self.min_bytes_input.update(cx, |input, cx| { + self.min_excerpt_bytes_input.update(cx, |input, cx| { input.set_text(options.excerpt.min_bytes.to_string(), window, cx); }); self.cursor_context_ratio_input.update(cx, |input, cx| { @@ -163,6 +167,9 @@ impl Zeta2Inspector { cx, ); }); + self.max_prompt_bytes_input.update(cx, |input, cx| { + input.set_text(options.max_prompt_bytes.to_string(), window, cx); + }); cx.notify(); } @@ -236,8 +243,8 @@ impl Zeta2Inspector { } let excerpt_options = EditPredictionExcerptOptions { - max_bytes: number_input_value(&this.max_bytes_input, cx), - min_bytes: number_input_value(&this.min_bytes_input, cx), + max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx), + min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx), target_before_cursor_over_total_bytes: number_input_value( &this.cursor_context_ratio_input, cx, @@ -247,7 +254,8 @@ impl Zeta2Inspector { this.set_options( ZetaOptions { excerpt: excerpt_options, - ..this.zeta.read(cx).options().clone() + max_prompt_bytes: number_input_value(&this.max_prompt_bytes_input, cx), + max_diagnostic_bytes: this.zeta.read(cx).options().max_diagnostic_bytes, }, cx, ); @@ -520,16 +528,15 @@ impl Render for Zeta2Inspector { .child( v_flex() .gap_2() - .child( - Headline::new("Excerpt Options").size(HeadlineSize::Small), - ) + .child(Headline::new("Options").size(HeadlineSize::Small)) .child( h_flex() .gap_2() .items_end() - .child(self.max_bytes_input.clone()) - .child(self.min_bytes_input.clone()) + .child(self.max_excerpt_bytes_input.clone()) + .child(self.min_excerpt_bytes_input.clone()) .child(self.cursor_context_ratio_input.clone()) + .child(self.max_prompt_bytes_input.clone()) .child( ui::Button::new("reset-options", "Reset") .disabled( diff --git a/crates/zeta_cli/src/main.rs b/crates/zeta_cli/src/main.rs index 44ed567a5df1dce4862523d0e2f1c04b44cffa94..40380c9ae143f6588c4a23565a25ebd17b6f64c9 100644 --- a/crates/zeta_cli/src/main.rs +++ b/crates/zeta_cli/src/main.rs @@ -63,11 +63,11 @@ struct ContextArgs { #[derive(Debug, Args)] struct Zeta2Args { #[arg(long, default_value_t = 8192)] - prompt_max_bytes: usize, + max_prompt_bytes: usize, #[arg(long, default_value_t = 2048)] - excerpt_max_bytes: usize, + max_excerpt_bytes: usize, #[arg(long, default_value_t = 1024)] - excerpt_min_bytes: usize, + min_excerpt_bytes: usize, #[arg(long, default_value_t = 0.66)] target_before_cursor_over_total_bytes: f32, #[arg(long, default_value_t = 1024)] @@ -225,12 +225,13 @@ async fn get_context( zeta.register_buffer(&buffer, &project, cx); zeta.set_options(zeta2::ZetaOptions { excerpt: EditPredictionExcerptOptions { - max_bytes: zeta2_args.excerpt_max_bytes, - min_bytes: zeta2_args.excerpt_min_bytes, + max_bytes: zeta2_args.max_excerpt_bytes, + min_bytes: zeta2_args.min_excerpt_bytes, target_before_cursor_over_total_bytes: zeta2_args .target_before_cursor_over_total_bytes, }, max_diagnostic_bytes: zeta2_args.max_diagnostic_bytes, + max_prompt_bytes: zeta2_args.max_prompt_bytes, }) }); // TODO: Actually wait for indexing. @@ -246,7 +247,7 @@ async fn get_context( let planned_prompt = cloud_zeta2_prompt::PlannedPrompt::populate( &request, &cloud_zeta2_prompt::PlanOptions { - max_bytes: zeta2_args.prompt_max_bytes, + max_bytes: zeta2_args.max_prompt_bytes, }, )?; anyhow::Ok(planned_prompt.to_prompt_string())