History

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

Release Notes: - N/A

Agus Zubiaga committed 1 year ago

Commit 831de8e1 parent bc528418 files · +118 −33Browse files

Unified · Split

8 files
  1. Cargo.lock +2 −0
  2. crates/cloud_llm_client/src/predict_edits_v3.rs +2 −0
  3. crates/cloud_zeta2_prompt/Cargo.toml +1 −0
  4. crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs +77 −12
  5. crates/zeta2/Cargo.toml +1 −0
  6. crates/zeta2/src/zeta2.rs +7 −1
  7. crates/zeta2_tools/src/zeta2_tools.rs +21 −14
  8. crates/zeta_cli/src/main.rs +7 −6
Cargo.lock+2−0
@@ -3233,6 +3233,7 @@ version = "0.1.0" · 3232 unmodified lines
32333233 dependencies = [
32343234 "anyhow",
32353235 "cloud_llm_client",
3236+ "indoc",
32363237 "ordered-float 2.10.1",
32373238 "rustc-hash 2.1.1",
32383239 "strum 0.27.1",
@@ -21651,6 +21652,7 @@ dependencies = [ · 18412 unmodified lines
2165121652 "chrono",
2165221653 "client",
2165321654 "cloud_llm_client",
21655+ "cloud_zeta2_prompt",
2165421656 "edit_prediction",
2165521657 "edit_prediction_context",
2165621658 "futures 0.3.31",
crates/cloud_llm_client/src/predict_edits_v3.rs+2−0
@@ -29,8 +29,10 @@ pub struct PredictEditsRequest { · 28 unmodified lines
2929 /// Info about the git repository state, only present when can_collect_data is true.
3030 #[serde(skip_serializing_if = "Option::is_none", default)]
3131 pub git_info: Option<PredictEditsGitInfo>,
32+ // Only available to staff
3233 #[serde(default)]
3334 pub debug_info: bool,
35+ pub prompt_max_bytes: Option<usize>,
3436 }
3537
3638 #[derive(Debug, Clone, Serialize, Deserialize)]
crates/cloud_zeta2_prompt/Cargo.toml+1−0
@@ -14,6 +14,7 @@ path = "src/cloud_zeta2_prompt.rs" · 13 unmodified lines
1414 [dependencies]
1515 anyhow.workspace = true
1616 cloud_llm_client.workspace = true
17+indoc.workspace = true
1718 ordered-float.workspace = true
1819 rustc-hash.workspace = true
1920 strum.workspace = true
crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs+77−12
@@ -1,17 +1,28 @@
11 //! Zeta2 prompt planning and generation code shared with cloud.
22
33 use anyhow::{Result, anyhow};
4−use cloud_llm_client::predict_edits_v3::{self, ReferencedDeclaration};
4+use cloud_llm_client::predict_edits_v3::{self, Event, ReferencedDeclaration};
5+use indoc::indoc;
56 use ordered_float::OrderedFloat;
67 use rustc_hash::{FxHashMap, FxHashSet};
8+use std::fmt::Write;
79 use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path};
810 use strum::{EnumIter, IntoEnumIterator};
911
12+pub const DEFAULT_MAX_PROMPT_BYTES: usize = 10 * 1024;
13+
1014 pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
1115 /// NOTE: Differs from zed version of constant - includes a newline
12−pub const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>\n";
16+pub const EDITABLE_REGION_START_MARKER_WITH_NEWLINE: &str = "<|editable_region_start|>\n";
1317 /// NOTE: Differs from zed version of constant - includes a newline
14−pub const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>\n";
18+pub const EDITABLE_REGION_END_MARKER_WITH_NEWLINE: &str = "<|editable_region_end|>\n";
19+
20+// TODO: use constants for markers?
21+pub const SYSTEM_PROMPT: &str = indoc! {"
22+ 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.
23+
24+ 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.
25+"};
1526
1627 pub struct PlannedPrompt<'a> {
1728 request: &'a predict_edits_v3::PredictEditsRequest,
@@ -286,7 +297,7 @@ impl<'a> PlannedPrompt<'a> { · 268 unmodified lines
286297 let mut excerpt_file_insertions = vec![
287298 (
288299 self.request.excerpt_range.start,
289− EDITABLE_REGION_START_MARKER,
300+ EDITABLE_REGION_START_MARKER_WITH_NEWLINE,
290301 ),
291302 (
292303 self.request.excerpt_range.start + self.request.cursor_offset,
@@ -298,10 +309,67 @@ impl<'a> PlannedPrompt<'a> { · 5 unmodified lines
298309 .end
299310 .saturating_sub(0)
300311 .max(self.request.excerpt_range.start),
301− EDITABLE_REGION_END_MARKER,
312+ EDITABLE_REGION_END_MARKER_WITH_NEWLINE,
302313 ),
303314 ];
304315
316+ let mut output = String::new();
317+ output.push_str("## User Edits\n\n");
318+ Self::push_events(&mut output, &self.request.events);
319+
320+ output.push_str("\n## Code\n\n");
321+ Self::push_file_snippets(&mut output, &mut excerpt_file_insertions, file_snippets);
322+ output
323+ }
324+
325+ fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
326+ for event in events {
327+ match event {
328+ Event::BufferChange {
329+ path,
330+ old_path,
331+ diff,
332+ predicted,
333+ } => {
334+ if let Some(old_path) = &old_path
335+ && let Some(new_path) = &path
336+ {
337+ if old_path != new_path {
338+ writeln!(
339+ output,
340+ "User renamed {} to {}\n\n",
341+ old_path.display(),
342+ new_path.display()
343+ )
344+ .unwrap();
345+ }
346+ }
347+
348+ let path = path
349+ .as_ref()
350+ .map_or_else(|| "untitled".to_string(), |path| path.display().to_string());
351+
352+ if *predicted {
353+ writeln!(
354+ output,
355+ "User accepted prediction {:?}:\n```diff\n{}\n```\n",
356+ path, diff
357+ )
358+ .unwrap();
359+ } else {
360+ writeln!(output, "User edited {:?}:\n```diff\n{}\n```\n", path, diff)
361+ .unwrap();
362+ }
363+ }
364+ }
365+ }
366+ }
367+
368+ fn push_file_snippets(
369+ output: &mut String,
370+ excerpt_file_insertions: &mut Vec<(usize, &'static str)>,
371+ file_snippets: Vec<(&Path, Vec<&PlannedSnippet>, bool)>,
372+ ) {
305373 fn push_excerpt_file_range(
306374 range: Range<usize>,
307375 text: &str,
@@ -325,7 +393,6 @@ impl<'a> PlannedPrompt<'a> { · 17 unmodified lines
325393 output.push_str(&text[last_offset - range.start..]);
326394 }
327395
328− let mut output = String::new();
329396 for (file_path, mut snippets, is_excerpt_file) in file_snippets {
330397 output.push_str(&format!("```{}\n", file_path.display()));
331398
@@ -345,8 +412,8 @@ impl<'a> PlannedPrompt<'a> { · 13 unmodified lines
345412 push_excerpt_file_range(
346413 last_range.end..snippet.range.end,
347414 text,
348− &mut excerpt_file_insertions,
349− &mut output,
415+ excerpt_file_insertions,
416+ output,
350417 );
351418 } else {
352419 output.push_str(text);
@@ -361,8 +428,8 @@ impl<'a> PlannedPrompt<'a> { · 8 unmodified lines
361428 push_excerpt_file_range(
362429 snippet.range.clone(),
363430 snippet.text,
364− &mut excerpt_file_insertions,
365− &mut output,
431+ excerpt_file_insertions,
432+ output,
366433 );
367434 } else {
368435 output.push_str(snippet.text);
@@ -372,8 +439,6 @@ impl<'a> PlannedPrompt<'a> { · 3 unmodified lines
372439
373440 output.push_str("```\n\n");
374441 }
375−
376− output
377442 }
378443 }
379444
crates/zeta2/Cargo.toml+1−0
@@ -17,6 +17,7 @@ arrayvec.workspace = true · 16 unmodified lines
1717 chrono.workspace = true
1818 client.workspace = true
1919 cloud_llm_client.workspace = true
20+cloud_zeta2_prompt.workspace = true
2021 edit_prediction.workspace = true
2122 edit_prediction_context.workspace = true
2223 futures.workspace = true
crates/zeta2/src/zeta2.rs+7−1
@@ -6,6 +6,7 @@ use cloud_llm_client::predict_edits_v3::{self, Signature}; · 5 unmodified lines
66 use cloud_llm_client::{
77 EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME, ZED_VERSION_HEADER_NAME,
88 };
9+use cloud_zeta2_prompt::DEFAULT_MAX_PROMPT_BYTES;
910 use edit_prediction::{DataCollectionState, Direction, EditPredictionProvider};
1011 use edit_prediction_context::{
1112 DeclarationId, EditPredictionContext, EditPredictionExcerptOptions, SyntaxIndex,
@@ -49,6 +50,7 @@ pub const DEFAULT_EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPrediction · 37 unmodified lines
4950
5051 pub const DEFAULT_OPTIONS: ZetaOptions = ZetaOptions {
5152 excerpt: DEFAULT_EXCERPT_OPTIONS,
53+ max_prompt_bytes: DEFAULT_MAX_PROMPT_BYTES,
5254 max_diagnostic_bytes: 2048,
5355 };
5456
@@ -71,6 +73,7 @@ pub struct Zeta { · 16 unmodified lines
7173 #[derive(Debug, Clone, PartialEq)]
7274 pub struct ZetaOptions {
7375 pub excerpt: EditPredictionExcerptOptions,
76+ pub max_prompt_bytes: usize,
7477 pub max_diagnostic_bytes: usize,
7578 }
7679
@@ -408,6 +411,7 @@ impl Zeta { · 331 unmodified lines
408411 debug_context.is_some(),
409412 &worktree_snapshots,
410413 index_state.as_deref(),
414+ Some(options.max_prompt_bytes),
411415 );
412416
413417 let retrieval_time = chrono::Utc::now() - before_retrieval;
@@ -702,6 +706,7 @@ impl Zeta { · 288 unmodified lines
702706 debug_info,
703707 &worktree_snapshots,
704708 index_state.as_deref(),
709+ Some(options.max_prompt_bytes),
705710 )
706711 })
707712 })
@@ -1062,6 +1067,7 @@ fn make_cloud_request( · 354 unmodified lines
10621067 debug_info: bool,
10631068 worktrees: &Vec<worktree::Snapshot>,
10641069 index_state: Option<&SyntaxIndexState>,
1070+ prompt_max_bytes: Option<usize>,
10651071 ) -> predict_edits_v3::PredictEditsRequest {
10661072 let mut signatures = Vec::new();
10671073 let mut declaration_to_signature_index = HashMap::default();
@@ -1132,9 +1138,9 @@ fn make_cloud_request( · 64 unmodified lines
11321138 can_collect_data,
11331139 diagnostic_groups,
11341140 diagnostic_groups_truncated,
1135−
11361141 git_info,
11371142 debug_info,
1143+ prompt_max_bytes,
11381144 }
11391145 }
11401146
crates/zeta2_tools/src/zeta2_tools.rs+21−14
@@ -57,13 +57,16 @@ pub fn init(cx: &mut App) { · 56 unmodified lines
5757 .detach();
5858 }
5959
60+// TODO show included diagnostics, and events
61+
6062 pub struct Zeta2Inspector {
6163 focus_handle: FocusHandle,
6264 project: Entity<Project>,
6365 last_prediction: Option<LastPredictionState>,
64− max_bytes_input: Entity<SingleLineInput>,
65− min_bytes_input: Entity<SingleLineInput>,
66+ max_excerpt_bytes_input: Entity<SingleLineInput>,
67+ min_excerpt_bytes_input: Entity<SingleLineInput>,
6668 cursor_context_ratio_input: Entity<SingleLineInput>,
69+ max_prompt_bytes_input: Entity<SingleLineInput>,
6770 active_view: ActiveView,
6871 zeta: Entity<Zeta>,
6972 _active_editor_subscription: Option<Subscription>,
@@ -129,9 +132,10 @@ impl Zeta2Inspector { · 59 unmodified lines
129132 project: project.clone(),
130133 last_prediction: None,
131134 active_view: ActiveView::Context,
132− max_bytes_input: Self::number_input("Max Bytes", window, cx),
133− min_bytes_input: Self::number_input("Min Bytes", window, cx),
135+ max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
136+ min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx),
134137 cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx),
138+ max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx),
135139 zeta: zeta.clone(),
136140 _active_editor_subscription: None,
137141 _update_state_task: Task::ready(()),
@@ -147,10 +151,10 @@ impl Zeta2Inspector { · 9 unmodified lines
147151 window: &mut Window,
148152 cx: &mut Context<Self>,
149153 ) {
150− self.max_bytes_input.update(cx, |input, cx| {
154+ self.max_excerpt_bytes_input.update(cx, |input, cx| {
151155 input.set_text(options.excerpt.max_bytes.to_string(), window, cx);
152156 });
153− self.min_bytes_input.update(cx, |input, cx| {
157+ self.min_excerpt_bytes_input.update(cx, |input, cx| {
154158 input.set_text(options.excerpt.min_bytes.to_string(), window, cx);
155159 });
156160 self.cursor_context_ratio_input.update(cx, |input, cx| {
@@ -163,6 +167,9 @@ impl Zeta2Inspector { · 6 unmodified lines
163167 cx,
164168 );
165169 });
170+ self.max_prompt_bytes_input.update(cx, |input, cx| {
171+ input.set_text(options.max_prompt_bytes.to_string(), window, cx);
172+ });
166173 cx.notify();
167174 }
168175
@@ -236,8 +243,8 @@ impl Zeta2Inspector { · 67 unmodified lines
236243 }
237244
238245 let excerpt_options = EditPredictionExcerptOptions {
239− max_bytes: number_input_value(&this.max_bytes_input, cx),
240− min_bytes: number_input_value(&this.min_bytes_input, cx),
246+ max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
247+ min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx),
241248 target_before_cursor_over_total_bytes: number_input_value(
242249 &this.cursor_context_ratio_input,
243250 cx,
@@ -247,7 +254,8 @@ impl Zeta2Inspector { · 3 unmodified lines
247254 this.set_options(
248255 ZetaOptions {
249256 excerpt: excerpt_options,
250− ..this.zeta.read(cx).options().clone()
257+ max_prompt_bytes: number_input_value(&this.max_prompt_bytes_input, cx),
258+ max_diagnostic_bytes: this.zeta.read(cx).options().max_diagnostic_bytes,
251259 },
252260 cx,
253261 );
@@ -520,16 +528,15 @@ impl Render for Zeta2Inspector { · 266 unmodified lines
520528 .child(
521529 v_flex()
522530 .gap_2()
523− .child(
524− Headline::new("Excerpt Options").size(HeadlineSize::Small),
525− )
531+ .child(Headline::new("Options").size(HeadlineSize::Small))
526532 .child(
527533 h_flex()
528534 .gap_2()
529535 .items_end()
530− .child(self.max_bytes_input.clone())
531− .child(self.min_bytes_input.clone())
536+ .child(self.max_excerpt_bytes_input.clone())
537+ .child(self.min_excerpt_bytes_input.clone())
532538 .child(self.cursor_context_ratio_input.clone())
539+ .child(self.max_prompt_bytes_input.clone())
533540 .child(
534541 ui::Button::new("reset-options", "Reset")
535542 .disabled(
crates/zeta_cli/src/main.rs+7−6
@@ -63,11 +63,11 @@ struct ContextArgs { · 62 unmodified lines
6363 #[derive(Debug, Args)]
6464 struct Zeta2Args {
6565 #[arg(long, default_value_t = 8192)]
66− prompt_max_bytes: usize,
66+ max_prompt_bytes: usize,
6767 #[arg(long, default_value_t = 2048)]
68− excerpt_max_bytes: usize,
68+ max_excerpt_bytes: usize,
6969 #[arg(long, default_value_t = 1024)]
70− excerpt_min_bytes: usize,
70+ min_excerpt_bytes: usize,
7171 #[arg(long, default_value_t = 0.66)]
7272 target_before_cursor_over_total_bytes: f32,
7373 #[arg(long, default_value_t = 1024)]
@@ -225,12 +225,13 @@ async fn get_context( · 151 unmodified lines
225225 zeta.register_buffer(&buffer, &project, cx);
226226 zeta.set_options(zeta2::ZetaOptions {
227227 excerpt: EditPredictionExcerptOptions {
228− max_bytes: zeta2_args.excerpt_max_bytes,
229− min_bytes: zeta2_args.excerpt_min_bytes,
228+ max_bytes: zeta2_args.max_excerpt_bytes,
229+ min_bytes: zeta2_args.min_excerpt_bytes,
230230 target_before_cursor_over_total_bytes: zeta2_args
231231 .target_before_cursor_over_total_bytes,
232232 },
233233 max_diagnostic_bytes: zeta2_args.max_diagnostic_bytes,
234+ max_prompt_bytes: zeta2_args.max_prompt_bytes,
234235 })
235236 });
236237 // TODO: Actually wait for indexing.
@@ -246,7 +247,7 @@ async fn get_context( · 9 unmodified lines
246247 let planned_prompt = cloud_zeta2_prompt::PlannedPrompt::populate(
247248 &request,
248249 &cloud_zeta2_prompt::PlanOptions {
249− max_bytes: zeta2_args.prompt_max_bytes,
250+ max_bytes: zeta2_args.max_prompt_bytes,
250251 },
251252 )?;
252253 anyhow::Ok(planned_prompt.to_prompt_string())