zeta2: Include edits in prompt and add max_prompt_bytes param (#38737)

Release Notes:

- N/A

Co-authored-by: Michael Sloan <mgsloan@gmail.com>

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",

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<PredictEditsGitInfo>,
+    // Only available to staff
     #[serde(default)]
     pub debug_info: bool,
+    pub prompt_max_bytes: Option<usize>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

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

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<usize>,
             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,
+                            excerpt_file_insertions,
-                            &mut output,
+                            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,
+                        excerpt_file_insertions,
-                        &mut output,
+                        output,
                     );
                 } else {
                     output.push_str(snippet.text);
@@ -372,8 +439,6 @@ impl<'a> PlannedPrompt<'a> {
 
             output.push_str("```\n\n");
         }
-
-        output
     }
 }
 

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

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<worktree::Snapshot>,
     index_state: Option<&SyntaxIndexState>,
+    prompt_max_bytes: Option<usize>,
 ) -> 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,
     }
 }
 

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<Project>,
     last_prediction: Option<LastPredictionState>,
-    max_bytes_input: Entity<SingleLineInput>,
+    max_excerpt_bytes_input: Entity<SingleLineInput>,
-    min_bytes_input: Entity<SingleLineInput>,
+    min_excerpt_bytes_input: Entity<SingleLineInput>,
     cursor_context_ratio_input: Entity<SingleLineInput>,
+    max_prompt_bytes_input: Entity<SingleLineInput>,
     active_view: ActiveView,
     zeta: Entity<Zeta>,
     _active_editor_subscription: Option<Subscription>,
@@ -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),
+            max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
-            min_bytes_input: Self::number_input("Min 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>,
     ) {
-        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),
+                    max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
-                    min_bytes: number_input_value(&this.min_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(
+                                    .child(Headline::new("Options").size(HeadlineSize::Small))
-                                        Headline::new("Excerpt Options").size(HeadlineSize::Small),
-                                    )
                                     .child(
                                         h_flex()
                                             .gap_2()
                                             .items_end()
-                                            .child(self.max_bytes_input.clone())
+                                            .child(self.max_excerpt_bytes_input.clone())
-                                            .child(self.min_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(

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,
+                            max_bytes: zeta2_args.max_excerpt_bytes,
-                            min_bytes: zeta2_args.excerpt_min_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())