zeta2: Support qwen3-minimal prompt format (#42902)

Oleksiy Syvokon created

This prompt is for a fine-tuned model. It has the following changes,
compared to `minimal`:
- No instructions at all, except for one sentence at the beginning of
the prompt.
- Output is a simplified unified diff -- hunk headers have no line
counts (e.g., `@@ -20 +20 @@`)
- Qwen's FIM tokens are used where possible (`<|file_sep|>`,
`<|fim_prefix|>`, `<|fim_suffix|>`, etc.)

To evaluate this model:
```
ZED_ZETA2_MODEL=zeta2-exp [usual zeta-cli eval params ...]  --prompt-format minimal-qwen
```

This will point to the most recent Baseten deployment of zeta2-exp
(which may change in the future, so the prompt-format may get out of
sync).

Release Notes:

- N/A

Change summary

crates/cloud_llm_client/src/predict_edits_v3.rs     |  3 
crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs | 89 ++++++++++++++
crates/zeta2/src/udiff.rs                           | 10 +
crates/zeta2/src/zeta2.rs                           |  2 
crates/zeta_cli/src/main.rs                         |  2 
5 files changed, 98 insertions(+), 8 deletions(-)

Detailed changes

crates/cloud_llm_client/src/predict_edits_v3.rs 🔗

@@ -78,6 +78,8 @@ pub enum PromptFormat {
     OnlySnippets,
     /// One-sentence instructions used in fine-tuned models
     Minimal,
+    /// One-sentence instructions + FIM-like template
+    MinimalQwen,
 }
 
 impl PromptFormat {
@@ -105,6 +107,7 @@ impl std::fmt::Display for PromptFormat {
             PromptFormat::NumLinesUniDiff => write!(f, "Numbered Lines / Unified Diff"),
             PromptFormat::OldTextNewText => write!(f, "Old Text / New Text"),
             PromptFormat::Minimal => write!(f, "Minimal"),
+            PromptFormat::MinimalQwen => write!(f, "Minimal + Qwen FIM"),
         }
     }
 }

crates/cloud_zeta2_prompt/src/cloud_zeta2_prompt.rs 🔗

@@ -3,7 +3,8 @@ pub mod retrieval_prompt;
 
 use anyhow::{Context as _, Result, anyhow};
 use cloud_llm_client::predict_edits_v3::{
-    self, DiffPathFmt, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
+    self, DiffPathFmt, Event, Excerpt, IncludedFile, Line, Point, PromptFormat,
+    ReferencedDeclaration,
 };
 use indoc::indoc;
 use ordered_float::OrderedFloat;
@@ -166,6 +167,21 @@ const OLD_TEXT_NEW_TEXT_REMINDER: &str = indoc! {r#"
 pub fn build_prompt(
     request: &predict_edits_v3::PredictEditsRequest,
 ) -> Result<(String, SectionLabels)> {
+    let mut section_labels = Default::default();
+
+    match request.prompt_format {
+        PromptFormat::MinimalQwen => {
+            let prompt = MinimalQwenPrompt {
+                events: request.events.clone(),
+                cursor_point: request.cursor_point,
+                cursor_path: request.excerpt_path.clone(),
+                included_files: request.included_files.clone(),
+            };
+            return Ok((prompt.render(), section_labels));
+        }
+        _ => (),
+    };
+
     let mut insertions = match request.prompt_format {
         PromptFormat::MarkedExcerpt => vec![
             (
@@ -191,6 +207,7 @@ pub fn build_prompt(
             vec![(request.cursor_point, CURSOR_MARKER)]
         }
         PromptFormat::OnlySnippets => vec![],
+        PromptFormat::MinimalQwen => unreachable!(),
     };
 
     let mut prompt = match request.prompt_format {
@@ -200,6 +217,7 @@ pub fn build_prompt(
         PromptFormat::OldTextNewText => XML_TAGS_INSTRUCTIONS.to_string(),
         PromptFormat::OnlySnippets => String::new(),
         PromptFormat::Minimal => STUDENT_MODEL_INSTRUCTIONS.to_string(),
+        PromptFormat::MinimalQwen => unreachable!(),
     };
 
     if request.events.is_empty() {
@@ -251,8 +269,6 @@ pub fn build_prompt(
     prompt.push_str(excerpts_preamble);
     prompt.push('\n');
 
-    let mut section_labels = Default::default();
-
     if !request.referenced_declarations.is_empty() || !request.signatures.is_empty() {
         let syntax_based_prompt = SyntaxBasedPrompt::populate(request)?;
         section_labels = syntax_based_prompt.write(&mut insertions, &mut prompt)?;
@@ -769,6 +785,7 @@ impl<'a> SyntaxBasedPrompt<'a> {
                             writeln!(output, "<|section_{}|>", section_index).ok();
                         }
                     }
+                    PromptFormat::MinimalQwen => unreachable!(),
                 }
 
                 let push_full_snippet = |output: &mut String| {
@@ -878,3 +895,69 @@ fn declaration_size(declaration: &ReferencedDeclaration, style: DeclarationStyle
         DeclarationStyle::Declaration => declaration.text.len(),
     }
 }
+
+struct MinimalQwenPrompt {
+    events: Vec<Event>,
+    cursor_point: Point,
+    cursor_path: Arc<Path>, // TODO: make a common struct with cursor_point
+    included_files: Vec<IncludedFile>,
+}
+
+impl MinimalQwenPrompt {
+    const INSTRUCTIONS: &str = "You are a code completion assistant that analyzes edit history to identify and systematically complete incomplete refactorings or patterns across the entire codebase.\n";
+
+    fn render(&self) -> String {
+        let edit_history = self.fmt_edit_history();
+        let context = self.fmt_context();
+
+        format!(
+            "{instructions}\n\n{edit_history}\n\n{context}",
+            instructions = MinimalQwenPrompt::INSTRUCTIONS,
+            edit_history = edit_history,
+            context = context
+        )
+    }
+
+    fn fmt_edit_history(&self) -> String {
+        if self.events.is_empty() {
+            "(No edit history)\n\n".to_string()
+        } else {
+            let mut events_str = String::new();
+            push_events(&mut events_str, &self.events);
+            format!(
+                "The following are the latest edits made by the user, from earlier to later.\n\n{}",
+                events_str
+            )
+        }
+    }
+
+    fn fmt_context(&self) -> String {
+        let mut context = String::new();
+        let include_line_numbers = true;
+
+        for related_file in &self.included_files {
+            writeln!(context, "<|file_sep|>{}", DiffPathFmt(&related_file.path)).unwrap();
+
+            if related_file.path == self.cursor_path {
+                write!(context, "<|fim_prefix|>").unwrap();
+                write_excerpts(
+                    &related_file.excerpts,
+                    &[(self.cursor_point, "<|fim_suffix|>")],
+                    related_file.max_row,
+                    include_line_numbers,
+                    &mut context,
+                );
+                writeln!(context, "<|fim_middle|>").unwrap();
+            } else {
+                write_excerpts(
+                    &related_file.excerpts,
+                    &[],
+                    related_file.max_row,
+                    include_line_numbers,
+                    &mut context,
+                );
+            }
+        }
+        context
+    }
+}

crates/zeta2/src/udiff.rs 🔗

@@ -391,10 +391,12 @@ impl<'a> DiffLine<'a> {
                 return Some(Self::HunkHeader(None));
             }
 
-            let (start_line_old, header) = header.strip_prefix('-')?.split_once(',')?;
-            let mut parts = header.split_ascii_whitespace();
-            let count_old = parts.next()?;
-            let (start_line_new, count_new) = parts.next()?.strip_prefix('+')?.split_once(',')?;
+            let mut tokens = header.split_whitespace();
+            let old_range = tokens.next()?.strip_prefix('-')?;
+            let new_range = tokens.next()?.strip_prefix('+')?;
+
+            let (start_line_old, count_old) = old_range.split_once(',').unwrap_or((old_range, "1"));
+            let (start_line_new, count_new) = new_range.split_once(',').unwrap_or((new_range, "1"));
 
             Some(Self::HunkHeader(Some(HunkLocation {
                 start_line_old: start_line_old.parse::<u32>().ok()?.saturating_sub(1),

crates/zeta2/src/zeta2.rs 🔗

@@ -1015,7 +1015,7 @@ impl Zeta {
                         // TODO: Implement parsing of multi-file diffs
                         crate::udiff::parse_diff(&output_text, get_buffer_from_context).await?
                     }
-                    PromptFormat::Minimal => {
+                    PromptFormat::Minimal | PromptFormat::MinimalQwen => {
                         if output_text.contains("--- a/\n+++ b/\nNo edits") {
                             let edits = vec![];
                             (&active_snapshot, edits)

crates/zeta_cli/src/main.rs 🔗

@@ -176,6 +176,7 @@ enum PromptFormat {
     NumberedLines,
     OldTextNewText,
     Minimal,
+    MinimalQwen,
 }
 
 impl Into<predict_edits_v3::PromptFormat> for PromptFormat {
@@ -187,6 +188,7 @@ impl Into<predict_edits_v3::PromptFormat> for PromptFormat {
             Self::NumberedLines => predict_edits_v3::PromptFormat::NumLinesUniDiff,
             Self::OldTextNewText => predict_edits_v3::PromptFormat::OldTextNewText,
             Self::Minimal => predict_edits_v3::PromptFormat::Minimal,
+            Self::MinimalQwen => predict_edits_v3::PromptFormat::MinimalQwen,
         }
     }
 }