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 b34c98fa72c..d2e287de62a 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 90f2a8b24fd..eeca7ed4e24 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 a1194f13615..b06431baf65 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 6690380c74b..89cfb4c41f2 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 d362441cdd5..11ca5fcfddd 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 4319f946f6c..47afc797d5d 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 de2cc660bef..2dfa292c438 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 44ed567a5df..40380c9ae14 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())