From aba82e90f1af613fe36f85f26c5a7cff590877f5 Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Tue, 28 Apr 2026 04:53:00 -0500 Subject: [PATCH] ep: Add diagnostics teacher format (#54999) Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes #ISSUE Release Notes: - N/A or Added/Fixed/Improved ... --- .../edit_prediction_cli/src/format_prompt.rs | 51 ++++++++- .../src/prompts/teacher.md | 1 + crates/zeta_prompt/src/zeta_prompt.rs | 106 ++++++++++++++++-- 3 files changed, 146 insertions(+), 12 deletions(-) diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs index 91f6aebe0f..64ca0585fa 100644 --- a/crates/edit_prediction_cli/src/format_prompt.rs +++ b/crates/edit_prediction_cli/src/format_prompt.rs @@ -47,7 +47,14 @@ pub async fn run_format_prompt( let (editable_range, context_range) = resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format); - let prompt = TeacherPrompt::format_prompt(example, editable_range, context_range); + let include_diagnostics = matches!(zeta_format, ZetaFormat::V0420Diagnostics); + + let prompt = TeacherPrompt::format_prompt( + example, + editable_range, + context_range, + include_diagnostics, + ); example.prompt = Some(ExamplePrompt { input: prompt, expected_output: None, @@ -64,8 +71,14 @@ pub async fn run_format_prompt( let (editable_range, context_range) = resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format); - let prompt = - TeacherMultiRegionPrompt::format_prompt(example, editable_range, context_range); + let include_diagnostics = matches!(zeta_format, ZetaFormat::V0420Diagnostics); + + let prompt = TeacherMultiRegionPrompt::format_prompt( + example, + editable_range, + context_range, + include_diagnostics, + ); example.prompt = Some(ExamplePrompt { input: prompt, expected_output: None, @@ -128,15 +141,20 @@ impl TeacherPrompt { example: &Example, editable_range: Range, context_range: Range, + include_diagnostics: bool, ) -> String { let edit_history = Self::format_edit_history(&example.spec.edit_history); let context = Self::format_context(example); let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range); + let diagnostics = include_diagnostics + .then(|| Self::format_diagnostics(example)) + .map(|diagnostics| format!("# 4. Diagnostics\n\n{diagnostics}")); let prompt_template = crate::prompt_assets::get_prompt("teacher.md"); let prompt = prompt_template .replace("{{context}}", &context) .replace("{{edit_history}}", &edit_history) + .replace("{{diagnostics}}", diagnostics.as_deref().unwrap_or("")) .replace("{{cursor_excerpt}}", &cursor_excerpt); prompt @@ -294,6 +312,27 @@ impl TeacherPrompt { let region = &text[start..end]; Ok(region.strip_suffix('\n').unwrap_or(region).to_string()) } + + fn format_diagnostics(example: &Example) -> String { + example + .prompt_inputs + .as_ref() + .map(|prompt_inputs| { + prompt_inputs + .active_buffer_diagnostics + .iter() + .map(|diagnostic| { + format!( + "*{}*:\n```\n{}\n```\n", + &diagnostic.message, &diagnostic.snippet + ) + }) + .collect::>() + .join("\n") + }) + .filter(|m| !m.is_empty()) + .unwrap_or("No Diagnostics".to_string()) + } } pub struct TeacherMultiRegionPrompt; @@ -309,15 +348,20 @@ impl TeacherMultiRegionPrompt { example: &Example, editable_range: Range, context_range: Range, + include_diagnostics: bool, ) -> String { let edit_history = Self::format_edit_history(&example.spec.edit_history); let context = Self::format_context(example); let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range); + let diagnostics = include_diagnostics + .then(|| TeacherPrompt::format_diagnostics(example)) + .map(|diagnostics| format!("# 4. Diagnostics\n\n{diagnostics}")); let prompt_template = crate::prompt_assets::get_prompt("teacher_multi_region.md"); let prompt = prompt_template .replace("{{context}}", &context) .replace("{{edit_history}}", &edit_history) + .replace("{{diagnostics}}", diagnostics.as_deref().unwrap_or("")) .replace("{{cursor_excerpt}}", &cursor_excerpt); prompt @@ -900,6 +944,7 @@ mod tests { }, editable_range, context_range, + false, ); assert!(prompt.contains(TeacherPrompt::EDITABLE_REGION_START)); diff --git a/crates/edit_prediction_cli/src/prompts/teacher.md b/crates/edit_prediction_cli/src/prompts/teacher.md index 524109c7a4..22a2514974 100644 --- a/crates/edit_prediction_cli/src/prompts/teacher.md +++ b/crates/edit_prediction_cli/src/prompts/teacher.md @@ -350,6 +350,7 @@ def calculate_square_perimeter(side): {{cursor_excerpt}} +{{diagnostics}} ----- diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 37799d528b..12767122fc 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -95,6 +95,8 @@ pub enum ZetaFormat { V0318SeedMultiRegions, /// V0318-style markers over the full available current file excerpt with no related files. V0327SingleFile, + /// V0318-style prompt with buffer diagnostics + V0420Diagnostics, } impl std::fmt::Display for ZetaFormat { @@ -245,7 +247,8 @@ pub fn format_zeta_prompt(input: &ZetaPromptInput, format: ZetaFormat) -> Option | ZetaFormat::V0316SeedMultiRegions | ZetaFormat::V0317SeedMultiRegions | ZetaFormat::V0331SeedCoderModelPy - | ZetaFormat::V0318SeedMultiRegions => 4096, + | ZetaFormat::V0318SeedMultiRegions + | ZetaFormat::V0420Diagnostics => 4096, ZetaFormat::V0327SingleFile => 16384, }; @@ -278,7 +281,7 @@ pub fn special_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] ]; TOKENS } - ZetaFormat::V0318SeedMultiRegions => { + ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => { static TOKENS: &[&str] = &[ seed_coder::FIM_SUFFIX, seed_coder::FIM_PREFIX, @@ -345,6 +348,7 @@ pub fn token_limits_for_format(format: ZetaFormat) -> (usize, usize) { | ZetaFormat::V0306SeedMultiRegions | ZetaFormat::V0316SeedMultiRegions | ZetaFormat::V0318SeedMultiRegions + | ZetaFormat::V0420Diagnostics | ZetaFormat::V0317SeedMultiRegions | ZetaFormat::V0327SingleFile | ZetaFormat::V0304SeedNoEdits => (350, 150), @@ -368,7 +372,9 @@ pub fn stop_tokens_for_format(format: ZetaFormat) -> &'static [&'static str] { | ZetaFormat::V0306SeedMultiRegions | ZetaFormat::V0304SeedNoEdits => &[], ZetaFormat::V0316SeedMultiRegions => &[multi_region::V0316_END_MARKER], - ZetaFormat::V0318SeedMultiRegions => &[multi_region::V0318_END_MARKER], + ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => { + &[multi_region::V0318_END_MARKER] + } ZetaFormat::V0317SeedMultiRegions => &[multi_region::V0317_END_MARKER], ZetaFormat::V0327SingleFile => &[multi_region::V0327_END_MARKER], } @@ -398,7 +404,8 @@ pub fn excerpt_ranges_for_format( | ZetaFormat::V0306SeedMultiRegions | ZetaFormat::V0316SeedMultiRegions | ZetaFormat::V0318SeedMultiRegions - | ZetaFormat::V0317SeedMultiRegions => ( + | ZetaFormat::V0317SeedMultiRegions + | ZetaFormat::V0420Diagnostics => ( ranges.editable_350.clone(), ranges.editable_350_context_150.clone(), ), @@ -497,7 +504,7 @@ pub fn write_cursor_excerpt_section_for_format( cursor_offset, )); } - ZetaFormat::V0318SeedMultiRegions => { + ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => { prompt.push_str(&build_v0318_cursor_prefix( path, context, @@ -709,7 +716,8 @@ pub fn format_prompt_with_budget_for_format( | ZetaFormat::V0306SeedMultiRegions | ZetaFormat::V0316SeedMultiRegions | ZetaFormat::V0318SeedMultiRegions - | ZetaFormat::V0317SeedMultiRegions => { + | ZetaFormat::V0317SeedMultiRegions + | ZetaFormat::V0420Diagnostics => { let mut cursor_section = String::new(); write_cursor_excerpt_section_for_format( format, @@ -720,6 +728,10 @@ pub fn format_prompt_with_budget_for_format( cursor_offset, ); + if format == ZetaFormat::V0420Diagnostics { + cursor_section.push_str(&format_active_buffer_diagnostics(input)); + } + let budget_with_margin = apply_prompt_budget_margin(max_tokens); seed_coder::assemble_fim_prompt( context, @@ -795,6 +807,26 @@ pub fn format_prompt_with_budget_for_format( return Some(prompt); } +fn format_active_buffer_diagnostics(input: &ZetaPromptInput) -> String { + let mut output = format!("{}diagnostics\n", seed_coder::FILE_MARKER); + + if input.active_buffer_diagnostics.is_empty() { + output.push_str("No Diagnostics\n"); + return output; + } + + for diagnostic in &input.active_buffer_diagnostics { + writeln!( + output, + "*{}*:\n```\n{}\n```", + diagnostic.message, diagnostic.snippet + ) + .ok(); + } + + output +} + pub fn filter_redundant_excerpts( mut related_files: Vec, cursor_path: &Path, @@ -829,6 +861,7 @@ pub fn max_edit_event_count_for_format(format: &ZetaFormat) -> usize { | ZetaFormat::V0316SeedMultiRegions | ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0317SeedMultiRegions + | ZetaFormat::V0420Diagnostics | ZetaFormat::V0327SingleFile => 6, } } @@ -854,6 +887,7 @@ pub fn get_prefill_for_format( | ZetaFormat::V0316SeedMultiRegions | ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0317SeedMultiRegions + | ZetaFormat::V0420Diagnostics | ZetaFormat::V0327SingleFile => String::new(), } } @@ -869,6 +903,7 @@ pub fn output_end_marker_for_format(format: ZetaFormat) -> Option<&'static str> | ZetaFormat::V0306SeedMultiRegions => Some(seed_coder::END_MARKER), ZetaFormat::V0316SeedMultiRegions => Some(multi_region::V0316_END_MARKER), ZetaFormat::V0318SeedMultiRegions => Some(multi_region::V0318_END_MARKER), + ZetaFormat::V0420Diagnostics => Some(multi_region::V0318_END_MARKER), ZetaFormat::V0317SeedMultiRegions => Some(multi_region::V0317_END_MARKER), ZetaFormat::V0327SingleFile => Some(multi_region::V0327_END_MARKER), @@ -914,7 +949,7 @@ pub fn encode_patch_as_output_for_format( Ok(None) } } - ZetaFormat::V0318SeedMultiRegions => { + ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => { let empty_patch = patch.lines().count() <= 3; if empty_patch { let marker_offsets = @@ -1001,7 +1036,7 @@ pub fn format_expected_output( multi_region::V0316_END_MARKER, ) } - ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0327SingleFile => { + ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => { let (new_editable, first_hunk_offset) = udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?; let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable); @@ -1013,6 +1048,18 @@ pub fn format_expected_output( multi_region::V0318_END_MARKER, ) } + ZetaFormat::V0327SingleFile => { + let (new_editable, first_hunk_offset) = + udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?; + let cursor_in_new = cursor_in_new_text(cursor_offset, first_hunk_offset, &new_editable); + multi_region::encode_from_old_and_new_v0318( + &old_editable, + &new_editable, + cursor_in_new, + CURSOR_MARKER, + multi_region::V0327_END_MARKER, + ) + } ZetaFormat::V0317SeedMultiRegions => { let (new_editable, first_hunk_offset) = udiff::apply_diff_to_string_with_hunk_offset(patch, &old_editable)?; @@ -1162,7 +1209,7 @@ pub fn parse_zeta2_model_output( editable_range_in_context, multi_region::apply_marker_span_v0316(old_editable_region, output)?, ), - ZetaFormat::V0318SeedMultiRegions => ( + ZetaFormat::V0318SeedMultiRegions | ZetaFormat::V0420Diagnostics => ( editable_range_in_context, multi_region::apply_marker_span_v0318(old_editable_region, output)?, ), @@ -5158,6 +5205,47 @@ mod tests { ); } + #[test] + fn test_v0420_formats_diagnostics_after_cursor_file() { + let mut input = make_input( + "prefix\neditable\nsuffix", + 7..15, + 10, + vec![], + vec![make_related_file("related.rs", "fn helper() {}\n")], + ); + input.active_buffer_diagnostics = vec![ActiveBufferDiagnostic { + severity: Some(1), + message: "missing semicolon".to_string(), + snippet: "let value = 1".to_string(), + snippet_buffer_row_range: 1..2, + diagnostic_range_in_snippet: 12..13, + }]; + + let prompt = + format_prompt_with_budget_for_format(&input, ZetaFormat::V0420Diagnostics, 10000) + .expect("v0420 prompt formatting should succeed"); + + assert_eq!( + prompt, + indoc! {r#" + <[fim-suffix]> + suffix + <[fim-prefix]>related.rs + fn helper() {} + + test.rs + prefix + <|marker_1|>edi<|user_cursor|>table<|marker_2|> + diagnostics + *missing semicolon*: + ``` + let value = 1 + ``` + <[fim-middle]>"#} + ); + } + #[test] fn test_v0317_formats_prompt_with_many_related_files() { let related_files = (0..900)