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
3233dependencies = [3233dependencies = [
3234 "anyhow",3234 "anyhow",
3235 "cloud_llm_client",3235 "cloud_llm_client",
3236 "indoc",
3236 "ordered-float 2.10.1",3237 "ordered-float 2.10.1",
3237 "rustc-hash 2.1.1",3238 "rustc-hash 2.1.1",
3238 "strum 0.27.1",3239 "strum 0.27.1",
@@ -21651,6 +21652,7 @@ dependencies = [ · 18412 unmodified lines
21651 "chrono",21652 "chrono",
21652 "client",21653 "client",
21653 "cloud_llm_client",21654 "cloud_llm_client",
21655 "cloud_zeta2_prompt",
21654 "edit_prediction",21656 "edit_prediction",
21655 "edit_prediction_context",21657 "edit_prediction_context",
21656 "futures 0.3.31",21658 "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
29 /// Info about the git repository state, only present when can_collect_data is true.29 /// Info about the git repository state, only present when can_collect_data is true.
30 #[serde(skip_serializing_if = "Option::is_none", default)]30 #[serde(skip_serializing_if = "Option::is_none", default)]
31 pub git_info: Option<PredictEditsGitInfo>,31 pub git_info: Option<PredictEditsGitInfo>,
32 // Only available to staff
32 #[serde(default)]33 #[serde(default)]
33 pub debug_info: bool,34 pub debug_info: bool,
35 pub prompt_max_bytes: Option<usize>,
34}36}
35 37
36#[derive(Debug, Clone, Serialize, Deserialize)]38#[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
14[dependencies]14[dependencies]
15anyhow.workspace = true15anyhow.workspace = true
16cloud_llm_client.workspace = true16cloud_llm_client.workspace = true
17indoc.workspace = true
17ordered-float.workspace = true18ordered-float.workspace = true
18rustc-hash.workspace = true19rustc-hash.workspace = true
19strum.workspace = true20strum.workspace = true
crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs+77−12
@@ -1,17 +1,28 @@
1//! Zeta2 prompt planning and generation code shared with cloud.1//! Zeta2 prompt planning and generation code shared with cloud.
2 2
3use anyhow::{Result, anyhow};3use anyhow::{Result, anyhow};
4use cloud_llm_client::predict_edits_v3::{self, ReferencedDeclaration};4use cloud_llm_client::predict_edits_v3::{self, Event, ReferencedDeclaration};
5use indoc::indoc;
5use ordered_float::OrderedFloat;6use ordered_float::OrderedFloat;
6use rustc_hash::{FxHashMap, FxHashSet};7use rustc_hash::{FxHashMap, FxHashSet};
8use std::fmt::Write;
7use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path};9use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path};
8use strum::{EnumIter, IntoEnumIterator};10use strum::{EnumIter, IntoEnumIterator};
9 11
12pub const DEFAULT_MAX_PROMPT_BYTES: usize = 10 * 1024;
13
10pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";14pub const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
11/// NOTE: Differs from zed version of constant - includes a newline15/// NOTE: Differs from zed version of constant - includes a newline
12pub const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>\n";16pub const EDITABLE_REGION_START_MARKER_WITH_NEWLINE: &str = "<|editable_region_start|>\n";
13/// NOTE: Differs from zed version of constant - includes a newline17/// NOTE: Differs from zed version of constant - includes a newline
14pub const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>\n";18pub const EDITABLE_REGION_END_MARKER_WITH_NEWLINE: &str = "<|editable_region_end|>\n";
19
20// TODO: use constants for markers?
21pub 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"};
15 26
16pub struct PlannedPrompt<'a> {27pub struct PlannedPrompt<'a> {
17 request: &'a predict_edits_v3::PredictEditsRequest,28 request: &'a predict_edits_v3::PredictEditsRequest,
@@ -286,7 +297,7 @@ impl<'a> PlannedPrompt<'a> { · 268 unmodified lines
286 let mut excerpt_file_insertions = vec![297 let mut excerpt_file_insertions = vec![
287 (298 (
288 self.request.excerpt_range.start,299 self.request.excerpt_range.start,
289 EDITABLE_REGION_START_MARKER,300 EDITABLE_REGION_START_MARKER_WITH_NEWLINE,
290 ),301 ),
291 (302 (
292 self.request.excerpt_range.start + self.request.cursor_offset,303 self.request.excerpt_range.start + self.request.cursor_offset,
@@ -298,10 +309,67 @@ impl<'a> PlannedPrompt<'a> { · 5 unmodified lines
298 .end309 .end
299 .saturating_sub(0)310 .saturating_sub(0)
300 .max(self.request.excerpt_range.start),311 .max(self.request.excerpt_range.start),
301 EDITABLE_REGION_END_MARKER,312 EDITABLE_REGION_END_MARKER_WITH_NEWLINE,
302 ),313 ),
303 ];314 ];
304 315
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 ) {
305 fn push_excerpt_file_range(373 fn push_excerpt_file_range(
306 range: Range<usize>,374 range: Range<usize>,
307 text: &str,375 text: &str,
@@ -325,7 +393,6 @@ impl<'a> PlannedPrompt<'a> { · 17 unmodified lines
325 output.push_str(&text[last_offset - range.start..]);393 output.push_str(&text[last_offset - range.start..]);
326 }394 }
327 395
328 let mut output = String::new();
329 for (file_path, mut snippets, is_excerpt_file) in file_snippets {396 for (file_path, mut snippets, is_excerpt_file) in file_snippets {
330 output.push_str(&format!("```{}\n", file_path.display()));397 output.push_str(&format!("```{}\n", file_path.display()));
331 398
@@ -345,8 +412,8 @@ impl<'a> PlannedPrompt<'a> { · 13 unmodified lines
345 push_excerpt_file_range(412 push_excerpt_file_range(
346 last_range.end..snippet.range.end,413 last_range.end..snippet.range.end,
347 text,414 text,
348 &mut excerpt_file_insertions,415 excerpt_file_insertions,
349 &mut output,416 output,
350 );417 );
351 } else {418 } else {
352 output.push_str(text);419 output.push_str(text);
@@ -361,8 +428,8 @@ impl<'a> PlannedPrompt<'a> { · 8 unmodified lines
361 push_excerpt_file_range(428 push_excerpt_file_range(
362 snippet.range.clone(),429 snippet.range.clone(),
363 snippet.text,430 snippet.text,
364 &mut excerpt_file_insertions,431 excerpt_file_insertions,
365 &mut output,432 output,
366 );433 );
367 } else {434 } else {
368 output.push_str(snippet.text);435 output.push_str(snippet.text);
@@ -372,8 +439,6 @@ impl<'a> PlannedPrompt<'a> { · 3 unmodified lines
372 439
373 output.push_str("```\n\n");440 output.push_str("```\n\n");
374 }441 }
375
376 output
377 }442 }
378}443}
379 444
crates/zeta2/Cargo.toml+1−0
@@ -17,6 +17,7 @@ arrayvec.workspace = true · 16 unmodified lines
17chrono.workspace = true17chrono.workspace = true
18client.workspace = true18client.workspace = true
19cloud_llm_client.workspace = true19cloud_llm_client.workspace = true
20cloud_zeta2_prompt.workspace = true
20edit_prediction.workspace = true21edit_prediction.workspace = true
21edit_prediction_context.workspace = true22edit_prediction_context.workspace = true
22futures.workspace = true23futures.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
6use cloud_llm_client::{6use cloud_llm_client::{
7 EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME, ZED_VERSION_HEADER_NAME,7 EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME, ZED_VERSION_HEADER_NAME,
8};8};
9use cloud_zeta2_prompt::DEFAULT_MAX_PROMPT_BYTES;
9use edit_prediction::{DataCollectionState, Direction, EditPredictionProvider};10use edit_prediction::{DataCollectionState, Direction, EditPredictionProvider};
10use edit_prediction_context::{11use edit_prediction_context::{
11 DeclarationId, EditPredictionContext, EditPredictionExcerptOptions, SyntaxIndex,12 DeclarationId, EditPredictionContext, EditPredictionExcerptOptions, SyntaxIndex,
@@ -49,6 +50,7 @@ pub const DEFAULT_EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPrediction · 37 unmodified lines
49 50
50pub const DEFAULT_OPTIONS: ZetaOptions = ZetaOptions {51pub const DEFAULT_OPTIONS: ZetaOptions = ZetaOptions {
51 excerpt: DEFAULT_EXCERPT_OPTIONS,52 excerpt: DEFAULT_EXCERPT_OPTIONS,
53 max_prompt_bytes: DEFAULT_MAX_PROMPT_BYTES,
52 max_diagnostic_bytes: 2048,54 max_diagnostic_bytes: 2048,
53};55};
54 56
@@ -71,6 +73,7 @@ pub struct Zeta { · 16 unmodified lines
71#[derive(Debug, Clone, PartialEq)]73#[derive(Debug, Clone, PartialEq)]
72pub struct ZetaOptions {74pub struct ZetaOptions {
73 pub excerpt: EditPredictionExcerptOptions,75 pub excerpt: EditPredictionExcerptOptions,
76 pub max_prompt_bytes: usize,
74 pub max_diagnostic_bytes: usize,77 pub max_diagnostic_bytes: usize,
75}78}
76 79
@@ -408,6 +411,7 @@ impl Zeta { · 331 unmodified lines
408 debug_context.is_some(),411 debug_context.is_some(),
409 &worktree_snapshots,412 &worktree_snapshots,
410 index_state.as_deref(),413 index_state.as_deref(),
414 Some(options.max_prompt_bytes),
411 );415 );
412 416
413 let retrieval_time = chrono::Utc::now() - before_retrieval;417 let retrieval_time = chrono::Utc::now() - before_retrieval;
@@ -702,6 +706,7 @@ impl Zeta { · 288 unmodified lines
702 debug_info,706 debug_info,
703 &worktree_snapshots,707 &worktree_snapshots,
704 index_state.as_deref(),708 index_state.as_deref(),
709 Some(options.max_prompt_bytes),
705 )710 )
706 })711 })
707 })712 })
@@ -1062,6 +1067,7 @@ fn make_cloud_request( · 354 unmodified lines
1062 debug_info: bool,1067 debug_info: bool,
1063 worktrees: &Vec<worktree::Snapshot>,1068 worktrees: &Vec<worktree::Snapshot>,
1064 index_state: Option<&SyntaxIndexState>,1069 index_state: Option<&SyntaxIndexState>,
1070 prompt_max_bytes: Option<usize>,
1065) -> predict_edits_v3::PredictEditsRequest {1071) -> predict_edits_v3::PredictEditsRequest {
1066 let mut signatures = Vec::new();1072 let mut signatures = Vec::new();
1067 let mut declaration_to_signature_index = HashMap::default();1073 let mut declaration_to_signature_index = HashMap::default();
@@ -1132,9 +1138,9 @@ fn make_cloud_request( · 64 unmodified lines
1132 can_collect_data,1138 can_collect_data,
1133 diagnostic_groups,1139 diagnostic_groups,
1134 diagnostic_groups_truncated,1140 diagnostic_groups_truncated,
1135
1136 git_info,1141 git_info,
1137 debug_info,1142 debug_info,
1143 prompt_max_bytes,
1138 }1144 }
1139}1145}
1140 1146
crates/zeta2_tools/src/zeta2_tools.rs+21−14
@@ -57,13 +57,16 @@ pub fn init(cx: &mut App) { · 56 unmodified lines
57 .detach();57 .detach();
58}58}
59 59
60// TODO show included diagnostics, and events
61
60pub struct Zeta2Inspector {62pub struct Zeta2Inspector {
61 focus_handle: FocusHandle,63 focus_handle: FocusHandle,
62 project: Entity<Project>,64 project: Entity<Project>,
63 last_prediction: Option<LastPredictionState>,65 last_prediction: Option<LastPredictionState>,
64 max_bytes_input: Entity<SingleLineInput>,66 max_excerpt_bytes_input: Entity<SingleLineInput>,
65 min_bytes_input: Entity<SingleLineInput>,67 min_excerpt_bytes_input: Entity<SingleLineInput>,
66 cursor_context_ratio_input: Entity<SingleLineInput>,68 cursor_context_ratio_input: Entity<SingleLineInput>,
69 max_prompt_bytes_input: Entity<SingleLineInput>,
67 active_view: ActiveView,70 active_view: ActiveView,
68 zeta: Entity<Zeta>,71 zeta: Entity<Zeta>,
69 _active_editor_subscription: Option<Subscription>,72 _active_editor_subscription: Option<Subscription>,
@@ -129,9 +132,10 @@ impl Zeta2Inspector { · 59 unmodified lines
129 project: project.clone(),132 project: project.clone(),
130 last_prediction: None,133 last_prediction: None,
131 active_view: ActiveView::Context,134 active_view: ActiveView::Context,
132 max_bytes_input: Self::number_input("Max Bytes", window, cx),135 max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
133 min_bytes_input: Self::number_input("Min Bytes", window, cx),136 min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx),
134 cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx),137 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),
135 zeta: zeta.clone(),139 zeta: zeta.clone(),
136 _active_editor_subscription: None,140 _active_editor_subscription: None,
137 _update_state_task: Task::ready(()),141 _update_state_task: Task::ready(()),
@@ -147,10 +151,10 @@ impl Zeta2Inspector { · 9 unmodified lines
147 window: &mut Window,151 window: &mut Window,
148 cx: &mut Context<Self>,152 cx: &mut Context<Self>,
149 ) {153 ) {
150 self.max_bytes_input.update(cx, |input, cx| {154 self.max_excerpt_bytes_input.update(cx, |input, cx| {
151 input.set_text(options.excerpt.max_bytes.to_string(), window, cx);155 input.set_text(options.excerpt.max_bytes.to_string(), window, cx);
152 });156 });
153 self.min_bytes_input.update(cx, |input, cx| {157 self.min_excerpt_bytes_input.update(cx, |input, cx| {
154 input.set_text(options.excerpt.min_bytes.to_string(), window, cx);158 input.set_text(options.excerpt.min_bytes.to_string(), window, cx);
155 });159 });
156 self.cursor_context_ratio_input.update(cx, |input, cx| {160 self.cursor_context_ratio_input.update(cx, |input, cx| {
@@ -163,6 +167,9 @@ impl Zeta2Inspector { · 6 unmodified lines
163 cx,167 cx,
164 );168 );
165 });169 });
170 self.max_prompt_bytes_input.update(cx, |input, cx| {
171 input.set_text(options.max_prompt_bytes.to_string(), window, cx);
172 });
166 cx.notify();173 cx.notify();
167 }174 }
168 175
@@ -236,8 +243,8 @@ impl Zeta2Inspector { · 67 unmodified lines
236 }243 }
237 244
238 let excerpt_options = EditPredictionExcerptOptions {245 let excerpt_options = EditPredictionExcerptOptions {
239 max_bytes: number_input_value(&this.max_bytes_input, cx),246 max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
240 min_bytes: number_input_value(&this.min_bytes_input, cx),247 min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx),
241 target_before_cursor_over_total_bytes: number_input_value(248 target_before_cursor_over_total_bytes: number_input_value(
242 &this.cursor_context_ratio_input,249 &this.cursor_context_ratio_input,
243 cx,250 cx,
@@ -247,7 +254,8 @@ impl Zeta2Inspector { · 3 unmodified lines
247 this.set_options(254 this.set_options(
248 ZetaOptions {255 ZetaOptions {
249 excerpt: excerpt_options,256 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,
251 },259 },
252 cx,260 cx,
253 );261 );
@@ -520,16 +528,15 @@ impl Render for Zeta2Inspector { · 266 unmodified lines
520 .child(528 .child(
521 v_flex()529 v_flex()
522 .gap_2()530 .gap_2()
523 .child(531 .child(Headline::new("Options").size(HeadlineSize::Small))
524 Headline::new("Excerpt Options").size(HeadlineSize::Small),
525 )
526 .child(532 .child(
527 h_flex()533 h_flex()
528 .gap_2()534 .gap_2()
529 .items_end()535 .items_end()
530 .child(self.max_bytes_input.clone())536 .child(self.max_excerpt_bytes_input.clone())
531 .child(self.min_bytes_input.clone())537 .child(self.min_excerpt_bytes_input.clone())
532 .child(self.cursor_context_ratio_input.clone())538 .child(self.cursor_context_ratio_input.clone())
539 .child(self.max_prompt_bytes_input.clone())
533 .child(540 .child(
534 ui::Button::new("reset-options", "Reset")541 ui::Button::new("reset-options", "Reset")
535 .disabled(542 .disabled(
crates/zeta_cli/src/main.rs+7−6
@@ -63,11 +63,11 @@ struct ContextArgs { · 62 unmodified lines
63#[derive(Debug, Args)]63#[derive(Debug, Args)]
64struct Zeta2Args {64struct Zeta2Args {
65 #[arg(long, default_value_t = 8192)]65 #[arg(long, default_value_t = 8192)]
66 prompt_max_bytes: usize,66 max_prompt_bytes: usize,
67 #[arg(long, default_value_t = 2048)]67 #[arg(long, default_value_t = 2048)]
68 excerpt_max_bytes: usize,68 max_excerpt_bytes: usize,
69 #[arg(long, default_value_t = 1024)]69 #[arg(long, default_value_t = 1024)]
70 excerpt_min_bytes: usize,70 min_excerpt_bytes: usize,
71 #[arg(long, default_value_t = 0.66)]71 #[arg(long, default_value_t = 0.66)]
72 target_before_cursor_over_total_bytes: f32,72 target_before_cursor_over_total_bytes: f32,
73 #[arg(long, default_value_t = 1024)]73 #[arg(long, default_value_t = 1024)]
@@ -225,12 +225,13 @@ async fn get_context( · 151 unmodified lines
225 zeta.register_buffer(&buffer, &project, cx);225 zeta.register_buffer(&buffer, &project, cx);
226 zeta.set_options(zeta2::ZetaOptions {226 zeta.set_options(zeta2::ZetaOptions {
227 excerpt: EditPredictionExcerptOptions {227 excerpt: EditPredictionExcerptOptions {
228 max_bytes: zeta2_args.excerpt_max_bytes,228 max_bytes: zeta2_args.max_excerpt_bytes,
229 min_bytes: zeta2_args.excerpt_min_bytes,229 min_bytes: zeta2_args.min_excerpt_bytes,
230 target_before_cursor_over_total_bytes: zeta2_args230 target_before_cursor_over_total_bytes: zeta2_args
231 .target_before_cursor_over_total_bytes,231 .target_before_cursor_over_total_bytes,
232 },232 },
233 max_diagnostic_bytes: zeta2_args.max_diagnostic_bytes,233 max_diagnostic_bytes: zeta2_args.max_diagnostic_bytes,
234 max_prompt_bytes: zeta2_args.max_prompt_bytes,
234 })235 })
235 });236 });
236 // TODO: Actually wait for indexing.237 // TODO: Actually wait for indexing.
@@ -246,7 +247,7 @@ async fn get_context( · 9 unmodified lines
246 let planned_prompt = cloud_zeta2_prompt::PlannedPrompt::populate(247 let planned_prompt = cloud_zeta2_prompt::PlannedPrompt::populate(
247 &request,248 &request,
248 &cloud_zeta2_prompt::PlanOptions {249 &cloud_zeta2_prompt::PlanOptions {
249 max_bytes: zeta2_args.prompt_max_bytes,250 max_bytes: zeta2_args.max_prompt_bytes,
250 },251 },
251 )?;252 )?;
252 anyhow::Ok(planned_prompt.to_prompt_string())253 anyhow::Ok(planned_prompt.to_prompt_string())