ep: Use FIM-like prompt for zeta2 (#46657)

Release Notes:

- N/A

---------

Co-authored-by: Max Brunsfeld <maxbrunsfeld@gmail.com>
This commit is contained in:
Agus ZubiagaandMax Brunsfeld authored and GitHub committed 2026-01-13 00:16:24 +00:00
1 parent 8ab49c325f
commit 90ec58836e
4 files changed
+40 -77

No files matched your search

@@ -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();
+1 -1
View File
@@ -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");
+34 -71
View File
@@ -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]);
cb(prompt);
prompt.push_str("<|fim_suffix|>\n");
prompt.push_str(&input.cursor_excerpt[input.editable_range_in_excerpt.end..]);
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");
}