Agus Zubiaga and Max Brunsfeld created
Release Notes:
- N/A
---------
Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>crates/cloud_llm_client/src/predict_edits_v3.rs | 2
crates/edit_prediction/src/edit_prediction_tests.rs | 8
crates/edit_prediction/src/zeta2.rs | 2
crates/zeta_prompt/src/zeta_prompt.rs | 105 ++++----------
4 files changed, 40 insertions(+), 77 deletions(-)
@@ -219,7 +219,7 @@ impl Sub for Line {
pub struct RawCompletionRequest {
pub model: String,
pub prompt: String,
- pub max_tokens: Option<u32>,
+ pub max_tokens: u32,
pub temperature: Option<f32>,
pub stop: Vec<Cow<'static, str>>,
}
@@ -1330,12 +1330,12 @@ async fn test_rejections_flushing(cx: &mut TestAppContext) {
fn model_response(request: RawCompletionRequest, diff_to_apply: &str) -> RawCompletionResponse {
let prompt = &request.prompt;
- let open = "<editable_region>\n";
- let close = "</editable_region>";
+ let current_marker = "<|fim_middle|>current\n";
+ let updated_marker = "<|fim_middle|>updated\n";
let cursor = "<|user_cursor|>";
- let start_ix = open.len() + prompt.find(open).unwrap();
- let end_ix = start_ix + &prompt[start_ix..].find(close).unwrap();
+ let start_ix = current_marker.len() + prompt.find(current_marker).unwrap();
+ let end_ix = start_ix + &prompt[start_ix..].find(updated_marker).unwrap();
let excerpt = prompt[start_ix..end_ix].replace(cursor, "");
let new_excerpt = apply_diff_to_string(diff_to_apply, &excerpt).unwrap();
@@ -80,7 +80,7 @@ pub fn request_prediction_with_zeta2(
prompt,
temperature: None,
stop: vec![],
- max_tokens: None,
+ max_tokens: 1024,
};
log::trace!("Sending edit prediction request");
@@ -78,88 +78,51 @@ pub fn format_zeta_prompt(input: &ZetaPromptInput) -> String {
}
pub fn write_related_files(prompt: &mut String, related_files: &[RelatedFile]) {
- push_delimited(prompt, "related_files", &[], |prompt| {
- for file in related_files {
- let path_str = file.path.to_string_lossy();
- push_delimited(prompt, "related_file", &[("path", &path_str)], |prompt| {
- for excerpt in &file.excerpts {
- push_delimited(
- prompt,
- "related_excerpt",
- &[(
- "lines",
- &format!(
- "{}-{}",
- excerpt.row_range.start + 1,
- excerpt.row_range.end + 1
- ),
- )],
- |prompt| {
- prompt.push_str(&excerpt.text);
- prompt.push('\n');
- },
- );
- }
- });
+ for file in related_files {
+ let path_str = file.path.to_string_lossy();
+ write!(prompt, "<|file_sep|>{}\n", path_str).ok();
+ for excerpt in &file.excerpts {
+ prompt.push_str(&excerpt.text);
+ if !prompt.ends_with('\n') {
+ prompt.push('\n');
+ }
+ if excerpt.row_range.end < file.max_row {
+ prompt.push_str("...\n");
+ }
}
- });
+ }
}
fn write_edit_history_section(prompt: &mut String, input: &ZetaPromptInput) {
- push_delimited(prompt, "edit_history", &[], |prompt| {
- if input.events.is_empty() {
- prompt.push_str("(No edit history)");
- } else {
- for event in &input.events {
- write_event(prompt, event);
- }
- }
- });
+ prompt.push_str("<|file_sep|>edit history\n");
+ for event in &input.events {
+ write_event(prompt, event);
+ }
}
fn write_cursor_excerpt_section(prompt: &mut String, input: &ZetaPromptInput) {
- push_delimited(prompt, "cursor_excerpt", &[], |prompt| {
- let path_str = input.cursor_path.to_string_lossy();
- push_delimited(prompt, "file", &[("path", &path_str)], |prompt| {
- prompt.push_str(&input.cursor_excerpt[..input.editable_range_in_excerpt.start]);
- push_delimited(prompt, "editable_region", &[], |prompt| {
- prompt.push_str(
- &input.cursor_excerpt
- [input.editable_range_in_excerpt.start..input.cursor_offset_in_excerpt],
- );
- prompt.push_str(CURSOR_MARKER);
- prompt.push_str(
- &input.cursor_excerpt
- [input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
- );
- });
- prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
- });
- });
-}
+ let path_str = input.cursor_path.to_string_lossy();
+ write!(prompt, "<|file_sep|>{}\n", path_str).ok();
-fn push_delimited(
- prompt: &mut String,
- tag: &'static str,
- arguments: &[(&str, &str)],
- cb: impl FnOnce(&mut String),
-) {
- if !prompt.ends_with("\n") {
- prompt.push('\n');
- }
- prompt.push('<');
- prompt.push_str(tag);
- for (arg_name, arg_value) in arguments {
- write!(prompt, " {}=\"{}\"", arg_name, arg_value).ok();
- }
- prompt.push_str(">\n");
+ prompt.push_str("<|fim_prefix|>\n");
+ prompt.push_str(&input.cursor_excerpt[..input.editable_range_in_excerpt.start]);
+
+ prompt.push_str("<|fim_suffix|>\n");
+ prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
- cb(prompt);
+ prompt.push_str("<|fim_middle|>current\n");
+ prompt.push_str(
+ &input.cursor_excerpt
+ [input.editable_range_in_excerpt.start..input.cursor_offset_in_excerpt],
+ );
+ prompt.push_str(CURSOR_MARKER);
+ prompt.push_str(
+ &input.cursor_excerpt[input.cursor_offset_in_excerpt..input.editable_range_in_excerpt.end],
+ );
if !prompt.ends_with('\n') {
prompt.push('\n');
}
- prompt.push_str("</");
- prompt.push_str(tag);
- prompt.push_str(">\n");
+
+ prompt.push_str("<|fim_middle|>updated\n");
}