From 90ec58836e025954e079cdc4bb1c9f9a7e682a2e Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Mon, 12 Jan 2026 21:16:24 -0300 Subject: [PATCH] ep: Use FIM-like prompt for zeta2 (#46657) Release Notes: - N/A --------- Co-authored-by: Max Brunsfeld --- .../cloud_llm_client/src/predict_edits_v3.rs | 2 +- .../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(-) diff --git a/crates/cloud_llm_client/src/predict_edits_v3.rs b/crates/cloud_llm_client/src/predict_edits_v3.rs index 49300a8b7a42df169e2d450d5d835a0f8aa99776..af3052de18c55f7f38e617cb2012b1721d590562 100644 --- a/crates/cloud_llm_client/src/predict_edits_v3.rs +++ b/crates/cloud_llm_client/src/predict_edits_v3.rs @@ -219,7 +219,7 @@ impl Sub for Line { pub struct RawCompletionRequest { pub model: String, pub prompt: String, - pub max_tokens: Option, + pub max_tokens: u32, pub temperature: Option, pub stop: Vec>, } diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index e463f1b6ee30a48001493390ca9c06615e066ef8..3d13d2d7f8cec902fbde8560b7e251f1273e6d24 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -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 = "\n"; - let close = ""; + 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(); diff --git a/crates/edit_prediction/src/zeta2.rs b/crates/edit_prediction/src/zeta2.rs index 035bf663461a22462a1a56aaa5e21383edd30a58..b8fc7ce37fb5d793963b97dec6deb2855bf47ecc 100644 --- a/crates/edit_prediction/src/zeta2.rs +++ b/crates/edit_prediction/src/zeta2.rs @@ -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"); diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 21fbca1ae10b715d0c11a31dc9390aada03fa157..cb09431098a474086c8ec48c43cb29d264eeb83b 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -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("\n"); + + prompt.push_str("<|fim_middle|>updated\n"); }