yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ep: Add diagnostics teacher format (#54999)

Browse Source 
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 ...
This commit is contained in:
Ben Kunkle authored and GitHub committed 2026-04-28 09:53:00 +00:00
1 parent 03c8df3547
commit aba82e90f1
3 files changed
+146 -12
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/edit_prediction_cli/src/format_prompt.rs
+48 -3
View File
@@ -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<usize>,
context_range: Range<usize>,
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::<Vec<_>>()
.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<usize>,
context_range: Range<usize>,
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));
crates/edit_prediction_cli/src/prompts/teacher.md
+1
View File
@@ -350,6 +350,7 @@ def calculate_square_perimeter(side):
{{cursor_excerpt}}
{{diagnostics}}
-----
crates/zeta_prompt/src/zeta_prompt.rs
+97 -9
View File
@@ -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<RelatedFile>,
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]><filename>related.rs
fn helper() {}
<filename>test.rs
prefix
<|marker_1|>edi<|user_cursor|>table<|marker_2|>
<filename>diagnostics
*missing semicolon*:
```
let value = 1
```
<[fim-middle]>"#}
);
}
#[test]
fn test_v0317_formats_prompt_with_many_related_files() {
let related_files = (0..900)