1use crate::{
2 FormatPromptArgs, PredictionProvider,
3 example::{ActualCursor, Example, ExamplePrompt},
4 headless::EpAppState,
5 progress::{ExampleProgress, Step},
6 retrieve_context::run_context_retrieval,
7};
8use anyhow::{Context as _, Result, anyhow};
9use gpui::AsyncApp;
10use std::ops::Range;
11use std::sync::Arc;
12use zeta_prompt::{
13 ZetaFormat, format_expected_output, format_zeta_prompt, multi_region, resolve_cursor_region,
14};
15
16fn resolved_excerpt_ranges_for_format(
17 input: &zeta_prompt::ZetaPromptInput,
18 format: ZetaFormat,
19) -> (Range<usize>, Range<usize>) {
20 let (_, editable_range_in_context, context_range, _) = resolve_cursor_region(input, format);
21 let editable_range = (context_range.start + editable_range_in_context.start)
22 ..(context_range.start + editable_range_in_context.end);
23 (editable_range, context_range)
24}
25
26pub async fn run_format_prompt(
27 example: &mut Example,
28 args: &FormatPromptArgs,
29 app_state: Arc<EpAppState>,
30 example_progress: &ExampleProgress,
31 cx: AsyncApp,
32) -> Result<()> {
33 run_context_retrieval(example, app_state.clone(), example_progress, cx.clone()).await?;
34
35 let step_progress = example_progress.start(Step::FormatPrompt);
36
37 let prompt_inputs = example
38 .prompt_inputs
39 .as_ref()
40 .context("prompt_inputs must be set after context retrieval")?;
41
42 match args.provider {
43 PredictionProvider::Teacher(_, zeta_format)
44 | PredictionProvider::TeacherNonBatching(_, zeta_format) => {
45 step_progress.set_substatus("formatting teacher prompt");
46
47 let (editable_range, context_range) =
48 resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
49
50 let prompt = TeacherPrompt::format_prompt(example, editable_range, context_range);
51 example.prompt = Some(ExamplePrompt {
52 input: prompt,
53 expected_output: None,
54 rejected_output: None,
55 prefill: None,
56 provider: args.provider,
57 });
58 }
59 PredictionProvider::TeacherMultiRegion(_)
60 | PredictionProvider::TeacherMultiRegionNonBatching(_) => {
61 step_progress.set_substatus("formatting teacher multi-region prompt");
62
63 let zeta_format = ZetaFormat::default();
64 let (editable_range, context_range) =
65 resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
66
67 let prompt =
68 TeacherMultiRegionPrompt::format_prompt(example, editable_range, context_range);
69 example.prompt = Some(ExamplePrompt {
70 input: prompt,
71 expected_output: None,
72 rejected_output: None,
73 prefill: None,
74 provider: args.provider,
75 });
76 }
77 PredictionProvider::Zeta2(zeta_format) => {
78 step_progress.set_substatus("formatting zeta2 prompt");
79
80 let prompt = format_zeta_prompt(prompt_inputs, zeta_format);
81 let prefill = zeta_prompt::get_prefill(prompt_inputs, zeta_format);
82 let expected_output = example
83 .spec
84 .expected_patches_with_cursor_positions()
85 .into_iter()
86 .next()
87 .and_then(|(expected_patch, expected_cursor_offset)| {
88 format_expected_output(
89 prompt_inputs,
90 zeta_format,
91 &expected_patch,
92 expected_cursor_offset,
93 )
94 .ok()
95 });
96
97 let rejected_output = example.spec.rejected_patch.as_ref().and_then(|patch| {
98 format_expected_output(prompt_inputs, zeta_format, patch, None).ok()
99 });
100
101 example.prompt = prompt.map(|prompt| ExamplePrompt {
102 input: prompt,
103 expected_output,
104 rejected_output,
105 provider: args.provider,
106 prefill: Some(prefill),
107 });
108 }
109 _ => {
110 panic!("Cannot format prompt for {:?}", args.provider);
111 }
112 };
113 Ok(())
114}
115
116pub struct TeacherPrompt;
117
118impl TeacherPrompt {
119 pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";
120 pub(crate) const EDITABLE_REGION_END: &str = "\n<|editable_region_end|>";
121 pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";
122 pub(crate) const NO_EDITS: &str = "NO_EDITS";
123
124 /// Truncate edit history to this number of last lines
125 const MAX_HISTORY_LINES: usize = 128;
126
127 pub fn format_prompt(
128 example: &Example,
129 editable_range: Range<usize>,
130 context_range: Range<usize>,
131 ) -> String {
132 let edit_history = Self::format_edit_history(&example.spec.edit_history);
133 let context = Self::format_context(example);
134 let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range);
135
136 let prompt_template = crate::prompt_assets::get_prompt("teacher.md");
137 let prompt = prompt_template
138 .replace("{{context}}", &context)
139 .replace("{{edit_history}}", &edit_history)
140 .replace("{{cursor_excerpt}}", &cursor_excerpt);
141
142 prompt
143 }
144
145 pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {
146 // Check if the model indicated no edits are needed
147 let no_edits = (String::new(), None);
148 if let Some(last_codeblock) = extract_last_codeblock(&response) {
149 if last_codeblock.trim() == Self::NO_EDITS {
150 return Ok(no_edits);
151 }
152 }
153
154 if response
155 .trim_end_matches(&[' ', '\n', '`'])
156 .ends_with(Self::NO_EDITS)
157 {
158 return Ok(no_edits);
159 }
160
161 // Extract updated (new) editable region from the model response.
162 let new_editable_region = Self::extract_editable_region(&response)?;
163 let cursor_offset = new_editable_region.find(Self::USER_CURSOR_MARKER);
164 let mut new_editable_region = new_editable_region.replace(Self::USER_CURSOR_MARKER, "");
165 let old_editable_region = Self::extract_editable_region(
166 &example
167 .prompt
168 .as_ref()
169 .context("example prompt missing")?
170 .input,
171 )?
172 .replace(Self::USER_CURSOR_MARKER, "");
173
174 let prompt_inputs = example
175 .prompt_inputs
176 .as_ref()
177 .context("example is missing prompt inputs")?;
178
179 // Normalize leading newlines: if old starts with newline but new doesn't,
180 // prepend newline to new to preserve whitespace structure.
181 // This handles the case where the model drops the leading blank line.
182 if old_editable_region.starts_with('\n') && !new_editable_region.starts_with('\n') {
183 new_editable_region.insert(0, '\n');
184 }
185
186 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
187 let (editable_region_offset, _) = excerpt
188 .match_indices(&old_editable_region)
189 .min_by_key(|(index, _)| index.abs_diff(prompt_inputs.cursor_offset_in_excerpt))
190 .context("editable region not found in prompt content")?;
191 let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
192
193 let editable_region_lines = old_editable_region.lines().count() as u32;
194 let diff = language::unified_diff_with_context(
195 &old_editable_region,
196 &new_editable_region,
197 editable_region_start_line as u32,
198 editable_region_start_line as u32,
199 editable_region_lines,
200 );
201
202 let diff = indoc::formatdoc! {"
203 --- a/{path}
204 +++ b/{path}
205 {diff}",
206 path = example.spec.cursor_path.to_string_lossy(),
207 diff = diff,
208 };
209
210 let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
211 ActualCursor::from_editable_region(
212 &example.spec.cursor_path,
213 editable_region_cursor_offset,
214 &new_editable_region,
215 excerpt,
216 editable_region_offset,
217 editable_region_start_line,
218 )
219 });
220
221 Ok((diff, actual_cursor))
222 }
223
224 fn format_edit_history(edit_history: &str) -> String {
225 let lines: Vec<&str> = edit_history.lines().collect();
226
227 if lines.is_empty() {
228 return "(No edit history)".to_string();
229 }
230
231 if lines.len() > Self::MAX_HISTORY_LINES {
232 let truncated = lines[lines.len() - Self::MAX_HISTORY_LINES..].join("\n");
233 format!("{truncated}\n[...truncated...]")
234 } else {
235 lines.join("\n")
236 }
237 }
238
239 pub fn format_context(example: &Example) -> String {
240 let related_files = example
241 .prompt_inputs
242 .as_ref()
243 .and_then(|pi| pi.related_files.as_deref());
244
245 let Some(related_files) = related_files else {
246 return "(No context)".to_string();
247 };
248
249 if related_files.is_empty() {
250 return "(No context)".to_string();
251 }
252
253 let prefix = "`````";
254 let suffix = "`````\n\n";
255 let max_tokens = 1024;
256 zeta_prompt::format_related_files_within_budget(related_files, &prefix, &suffix, max_tokens)
257 }
258
259 fn format_cursor_excerpt(
260 example: &Example,
261 editable_range: Range<usize>,
262 context_range: Range<usize>,
263 ) -> String {
264 let mut result = String::new();
265
266 let prompt_inputs = example.prompt_inputs.as_ref().unwrap();
267 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
268 let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
269
270 let path_str = example.spec.cursor_path.to_string_lossy();
271 result.push_str(&format!("`````{path_str}\n"));
272 result.push_str(&excerpt[context_range.start..editable_range.start]);
273 result.push_str(Self::EDITABLE_REGION_START);
274 result.push_str(&excerpt[editable_range.start..cursor_offset]);
275 result.push_str(Self::USER_CURSOR_MARKER);
276 result.push_str(&excerpt[cursor_offset..editable_range.end]);
277 result.push_str(Self::EDITABLE_REGION_END);
278 result.push_str(&excerpt[editable_range.end..context_range.end]);
279 result.push_str("\n`````");
280
281 result
282 }
283
284 pub fn extract_editable_region(text: &str) -> Result<String> {
285 let start = text
286 .rfind(Self::EDITABLE_REGION_START)
287 .map_or(0, |pos| pos + Self::EDITABLE_REGION_START.len());
288 let end = text.rfind(Self::EDITABLE_REGION_END).unwrap_or(text.len());
289
290 if start >= end {
291 return Err(anyhow!("Invalid editable region markers"));
292 }
293
294 let region = &text[start..end];
295 Ok(region.strip_suffix('\n').unwrap_or(region).to_string())
296 }
297}
298
299pub struct TeacherMultiRegionPrompt;
300
301impl TeacherMultiRegionPrompt {
302 pub(crate) const USER_CURSOR_MARKER: &str = "<|user_cursor|>";
303 pub(crate) const NO_EDITS: &str = "NO_EDITS";
304
305 /// Truncate edit history to this number of last lines
306 const MAX_HISTORY_LINES: usize = 128;
307
308 pub fn format_prompt(
309 example: &Example,
310 editable_range: Range<usize>,
311 context_range: Range<usize>,
312 ) -> String {
313 let edit_history = Self::format_edit_history(&example.spec.edit_history);
314 let context = Self::format_context(example);
315 let cursor_excerpt = Self::format_cursor_excerpt(example, editable_range, context_range);
316
317 let prompt_template = crate::prompt_assets::get_prompt("teacher_multi_region.md");
318 let prompt = prompt_template
319 .replace("{{context}}", &context)
320 .replace("{{edit_history}}", &edit_history)
321 .replace("{{cursor_excerpt}}", &cursor_excerpt);
322
323 prompt
324 }
325
326 pub fn parse(example: &Example, response: &str) -> Result<(String, Option<ActualCursor>)> {
327 let no_edits = (String::new(), None);
328 if let Some(last_codeblock) = extract_last_codeblock(&response) {
329 if last_codeblock.trim() == Self::NO_EDITS {
330 return Ok(no_edits);
331 }
332 }
333
334 if response.trim().ends_with(Self::NO_EDITS) {
335 return Ok(no_edits);
336 }
337
338 let prompt_inputs = example
339 .prompt_inputs
340 .as_ref()
341 .context("example is missing prompt inputs")?;
342
343 let zeta_format = ZetaFormat::default();
344 let (editable_range, _) = resolved_excerpt_ranges_for_format(prompt_inputs, zeta_format);
345 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
346 let old_editable_region = &excerpt[editable_range.clone()];
347 let marker_offsets = multi_region::compute_marker_offsets(old_editable_region);
348
349 let codeblock =
350 extract_last_codeblock(&response).context("no codeblock found in model response")?;
351 let (start_num, end_num, raw_new_span) = multi_region::extract_marker_span(&codeblock)?;
352
353 let start_idx = start_num
354 .checked_sub(1)
355 .context("marker numbers are 1-indexed")?;
356 let end_idx = end_num
357 .checked_sub(1)
358 .context("marker numbers are 1-indexed")?;
359 let start_byte = *marker_offsets
360 .get(start_idx)
361 .context("start marker number out of range")?;
362 let end_byte = *marker_offsets
363 .get(end_idx)
364 .context("end marker number out of range")?;
365
366 if start_byte > end_byte {
367 return Err(anyhow!("start marker must come before end marker"));
368 }
369
370 let cursor_in_span = raw_new_span.find(Self::USER_CURSOR_MARKER);
371 let new_span = raw_new_span.replace(Self::USER_CURSOR_MARKER, "");
372
373 let old_span = &old_editable_region[start_byte..end_byte];
374 let mut new_span = new_span;
375 if old_span.ends_with('\n') && !new_span.ends_with('\n') && !new_span.is_empty() {
376 new_span.push('\n');
377 }
378 if !old_span.ends_with('\n') && new_span.ends_with('\n') {
379 new_span.pop();
380 }
381
382 let mut new_editable_region = String::new();
383 new_editable_region.push_str(&old_editable_region[..start_byte]);
384 new_editable_region.push_str(&new_span);
385 new_editable_region.push_str(&old_editable_region[end_byte..]);
386
387 let cursor_offset = cursor_in_span.map(|pos| start_byte + pos);
388
389 if old_editable_region.starts_with('\n') && !new_editable_region.starts_with('\n') {
390 new_editable_region.insert(0, '\n');
391 }
392
393 let editable_region_offset = editable_range.start;
394 let editable_region_start_line = excerpt[..editable_region_offset].matches('\n').count();
395
396 let editable_region_lines = old_editable_region.lines().count() as u32;
397 let diff = language::unified_diff_with_context(
398 old_editable_region,
399 &new_editable_region,
400 editable_region_start_line as u32,
401 editable_region_start_line as u32,
402 editable_region_lines,
403 );
404
405 let diff = indoc::formatdoc! {"
406 --- a/{path}
407 +++ b/{path}
408 {diff}",
409 path = example.spec.cursor_path.to_string_lossy(),
410 diff = diff,
411 };
412
413 let actual_cursor = cursor_offset.map(|editable_region_cursor_offset| {
414 ActualCursor::from_editable_region(
415 &example.spec.cursor_path,
416 editable_region_cursor_offset,
417 &new_editable_region,
418 excerpt,
419 editable_region_offset,
420 editable_region_start_line,
421 )
422 });
423
424 Ok((diff, actual_cursor))
425 }
426
427 fn format_edit_history(edit_history: &str) -> String {
428 let lines: Vec<&str> = edit_history.lines().collect();
429
430 if lines.is_empty() {
431 return "(No edit history)".to_string();
432 }
433
434 if lines.len() > Self::MAX_HISTORY_LINES {
435 let truncated = lines[lines.len() - Self::MAX_HISTORY_LINES..].join("\n");
436 format!("{truncated}\n[...truncated...]")
437 } else {
438 lines.join("\n")
439 }
440 }
441
442 pub fn format_context(example: &Example) -> String {
443 let related_files = example
444 .prompt_inputs
445 .as_ref()
446 .and_then(|pi| pi.related_files.as_deref());
447 let Some(related_files) = related_files else {
448 return "(No context)".to_string();
449 };
450
451 if related_files.is_empty() {
452 return "(No context)".to_string();
453 }
454
455 let prefix = "`````";
456 let suffix = "`````\n\n";
457 let max_tokens = 1024;
458 zeta_prompt::format_related_files_within_budget(related_files, &prefix, &suffix, max_tokens)
459 }
460
461 fn format_cursor_excerpt(
462 example: &Example,
463 editable_range: Range<usize>,
464 context_range: Range<usize>,
465 ) -> String {
466 let mut result = String::new();
467
468 let prompt_inputs = example.prompt_inputs.as_ref().unwrap();
469 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
470 let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
471
472 let editable_text = &excerpt[editable_range.clone()];
473 let cursor_in_editable = cursor_offset - editable_range.start;
474
475 let path_str = example.spec.cursor_path.to_string_lossy();
476 result.push_str(&format!("`````{path_str}\n"));
477
478 result.push_str(&excerpt[context_range.start..editable_range.start]);
479
480 multi_region::write_editable_with_markers(
481 &mut result,
482 editable_text,
483 cursor_in_editable,
484 Self::USER_CURSOR_MARKER,
485 );
486
487 result.push_str(&excerpt[editable_range.end..context_range.end]);
488 result.push_str("\n`````");
489
490 result
491 }
492}
493
494/// Extract the cursor excerpt from an example.
495/// First tries to extract from an existing prompt, then falls back to constructing from prompt_inputs.
496pub fn extract_cursor_excerpt_from_example(example: &Example) -> Option<String> {
497 // If we have the original prompt, extract the cursor excerpt from it
498 if let Some(prompt) = &example.prompt {
499 // Find "# 3. Current File" section and extract the content
500 if let Some(start) = prompt.input.find("# 3. Current File") {
501 let content_start = prompt.input[start..].find('`').map(|i| start + i)?;
502 let backtick_count = prompt.input[content_start..]
503 .chars()
504 .take_while(|&c| c == '`')
505 .count();
506 let content_start = content_start + backtick_count;
507
508 // Find the path line and skip it
509 let newline_pos = prompt.input[content_start..].find('\n')?;
510 let text_start = content_start + newline_pos + 1;
511
512 // Find the closing backticks
513 let closing_pattern = "`".repeat(backtick_count);
514 let text_end = prompt.input[text_start..].find(&closing_pattern)?;
515 let cursor_excerpt = &prompt.input[text_start..text_start + text_end];
516
517 let path_str = example.spec.cursor_path.to_string_lossy();
518 return Some(format!("`````{path_str}\n{cursor_excerpt}`````"));
519 }
520 }
521
522 // Fallback: construct from prompt_inputs if available
523 let prompt_inputs = example.prompt_inputs.as_ref()?;
524 let excerpt = prompt_inputs.cursor_excerpt.as_ref();
525 let cursor_offset = prompt_inputs.cursor_offset_in_excerpt;
526
527 // Simple fallback: just show content around cursor with markers
528 let path_str = example.spec.cursor_path.to_string_lossy();
529 let mut result = format!("`````{path_str}\n");
530 result.push_str(TeacherPrompt::EDITABLE_REGION_START);
531 result.push_str(&excerpt[..cursor_offset]);
532 result.push_str(TeacherPrompt::USER_CURSOR_MARKER);
533 result.push_str(&excerpt[cursor_offset..]);
534 result.push_str(TeacherPrompt::EDITABLE_REGION_END);
535 result.push_str("\n`````");
536
537 Some(result)
538}
539
540pub(crate) fn extract_last_codeblock(text: &str) -> Option<String> {
541 let lines: Vec<&str> = text.lines().collect();
542
543 // Search from the end for a closing fence (line containing only backticks, 3+)
544 let mut closing_line_idx = None;
545 let mut backtick_count = 0;
546
547 for i in (0..lines.len()).rev() {
548 let line = lines[i].trim();
549 if line.len() >= 3 && line.chars().all(|c| c == '`') {
550 closing_line_idx = Some(i);
551 backtick_count = line.len();
552 break;
553 }
554 }
555
556 let closing_idx = closing_line_idx?;
557
558 // Search backwards for matching opening fence
559 // Opening fence starts with same backtick count, possibly followed by language/metadata
560 let opening_pattern = "`".repeat(backtick_count);
561
562 for i in (0..closing_idx).rev() {
563 let line = lines[i];
564 if line.starts_with(&opening_pattern) {
565 // Ensure it's exactly the right number of backticks (not more)
566 let rest = &line[backtick_count..];
567 if rest.is_empty() || !rest.starts_with('`') {
568 // Found matching opening fence
569 // Extract content between opening and closing (exclusive)
570 if closing_idx > i + 1 {
571 let content = lines[i + 1..closing_idx].join("\n");
572 // Preserve trailing newline to match previous behavior
573 return Some(format!("{}\n", content));
574 } else {
575 // Empty block
576 return Some(String::new());
577 }
578 }
579 }
580 }
581
582 None
583}
584
585#[cfg(test)]
586mod tests {
587 use super::*;
588
589 #[test]
590 fn test_extract_last_code_block() {
591 let text = indoc::indoc! {"
592 Some thinking
593
594 ```
595 first block
596 ```
597
598 `````path='something' lines=1:2
599 last block
600 `````
601 "};
602 let last_block = extract_last_codeblock(text).unwrap();
603 assert_eq!(last_block, "last block\n");
604 }
605
606 #[test]
607 fn test_extract_codeblock_with_nested_fences() {
608 let text = indoc::indoc! {"
609 `````
610 content with ``` inline
611 and ```python nested
612 more content
613 `````
614 "};
615 let last_block = extract_last_codeblock(text).unwrap();
616 assert_eq!(
617 last_block,
618 "content with ``` inline\nand ```python nested\nmore content\n"
619 );
620 }
621
622 #[test]
623 fn test_extract_codeblock_ignores_inline_backticks() {
624 let text = indoc::indoc! {"
625 `````
626 here is some `code` with inline backticks
627 and here```more```stuff
628 `````
629 "};
630 let last_block = extract_last_codeblock(text).unwrap();
631 assert_eq!(
632 last_block,
633 "here is some `code` with inline backticks\nand here```more```stuff\n"
634 );
635 }
636
637 #[test]
638 fn test_extract_editable_region_old_format() {
639 let text = indoc::indoc! {"
640 some lines
641 are
642 here
643 <|editable_region_start|>
644 one
645 two three
646
647 <|editable_region_end|>
648 more
649 lines here
650 "};
651 let parsed = TeacherPrompt::extract_editable_region(text).unwrap();
652 assert_eq!(
653 parsed,
654 indoc::indoc! {"
655 one
656 two three"}
657 );
658 }
659
660 #[test]
661 fn test_extract_editable_region_marker_format() {
662 let text = indoc::indoc! {"
663 some context
664 <|marker_1|>
665 one
666 two three
667 <|marker_2|>
668 more context
669 "};
670 let parsed = multi_region::extract_editable_region_from_markers(text).unwrap();
671 assert_eq!(parsed, "one\ntwo three");
672 }
673
674 #[test]
675 fn test_extract_editable_region_multi_markers() {
676 let text = indoc::indoc! {"
677 prefix
678 <|marker_1|>
679 aaa
680 bbb
681 <|marker_2|>
682 ccc
683 ddd
684 <|marker_3|>
685 suffix
686 "};
687 let parsed = multi_region::extract_editable_region_from_markers(text).unwrap();
688 // Intermediate marker and its trailing \n are stripped
689 assert_eq!(parsed, "aaa\nbbb\nccc\nddd");
690 }
691
692 #[test]
693 fn test_extract_last_codeblock_nested_bibtex() {
694 let text = indoc::indoc! {r#"
695 Looking at the edit history, I can see that a Citation section was just added.
696
697 `````
698 ## Collaborations
699 Our mission is to create a 4D generative model.
700
701 ## Citation
702
703 If you found Unique3D helpful, please cite our report:
704 ```bibtex
705 @misc{wu2024unique3d,
706 title={Unique3D},
707 }
708 ```
709 `````
710 "#};
711 let last_block = extract_last_codeblock(text).unwrap();
712 assert_eq!(
713 last_block,
714 indoc::indoc! {r#"
715 ## Collaborations
716 Our mission is to create a 4D generative model.
717
718 ## Citation
719
720 If you found Unique3D helpful, please cite our report:
721 ```bibtex
722 @misc{wu2024unique3d,
723 title={Unique3D},
724 }
725 ```
726 "#}
727 );
728 }
729
730 #[test]
731 fn test_extract_editable_region_no_markers() {
732 let text = indoc::indoc! {"
733 one
734 two three"};
735 let parsed = TeacherPrompt::extract_editable_region(text).unwrap();
736 assert_eq!(
737 parsed,
738 indoc::indoc! {"
739 one
740 two three"}
741 );
742 }
743
744 #[test]
745 fn test_parse_no_edits_response() {
746 let response = indoc::indoc! {"
747 The code is already complete. There is no clear next edit to make.
748
749 `````
750 NO_EDITS
751 `````
752 "};
753 let codeblock = extract_last_codeblock(response).unwrap();
754 assert_eq!(codeblock.trim(), TeacherPrompt::NO_EDITS);
755 }
756
757 #[test]
758 fn test_extract_codeblock_no_valid_block() {
759 // Text with no code blocks should return None
760 let text = "Just some plain text without any code blocks";
761 assert!(extract_last_codeblock(text).is_none());
762
763 // Unclosed code block should return None
764 let text = indoc::indoc! {"
765 ```
766 unclosed block
767 "};
768 assert!(extract_last_codeblock(text).is_none());
769
770 // Analysis text with nested markdown but no proper outer block
771 let text = indoc::indoc! {"
772 # Analysis
773 Looking at this:
774 ```
775 some code
776 ```
777 But then more analysis without wrapping block
778 "};
779 // This should find the inner block
780 let result = extract_last_codeblock(text).unwrap();
781 assert_eq!(result, "some code\n");
782 }
783
784 #[test]
785 fn test_extract_codeblock_no_trailing_newline() {
786 // Text ending without trailing newline after closing fence
787 let text = "`````\ncontent here\n`````";
788 let result = extract_last_codeblock(text).unwrap();
789 assert_eq!(result, "content here\n");
790 }
791
792 #[test]
793 fn test_parse_no_edits_response_with_trailing_backticks() {
794 let response = "NO_EDITS```";
795
796 let parsed = TeacherPrompt::parse(
797 &Example {
798 spec: edit_prediction::example_spec::ExampleSpec {
799 name: "test".to_string(),
800 repository_url: "https://github.com/zed-industries/zed.git".to_string(),
801 revision: "HEAD".to_string(),
802 tags: Vec::new(),
803 reasoning: None,
804 uncommitted_diff: String::new(),
805 cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
806 cursor_position: "0:0".to_string(),
807 edit_history: String::new(),
808 expected_patches: Vec::new(),
809 rejected_patch: None,
810 telemetry: None,
811 human_feedback: Vec::new(),
812 rating: None,
813 },
814 prompt_inputs: None,
815 prompt: None,
816 predictions: Vec::new(),
817 score: Vec::new(),
818 qa: Vec::new(),
819 zed_version: None,
820 state: None,
821 },
822 response,
823 )
824 .unwrap();
825
826 assert!(parsed.0.is_empty());
827 assert!(parsed.1.is_none());
828 }
829
830 #[test]
831 fn test_v0327_teacher_prompt_uses_resolved_ranges() {
832 let excerpt = (0..80)
833 .map(|index| format!("line{index:02}\n"))
834 .collect::<String>();
835 let cursor_offset = excerpt.find("line40").expect("cursor line exists");
836 let prompt_inputs = zeta_prompt::ZetaPromptInput {
837 cursor_path: std::path::Path::new("src/main.rs").into(),
838 cursor_excerpt: excerpt.clone().into(),
839 cursor_offset_in_excerpt: cursor_offset,
840 excerpt_start_row: None,
841 events: Vec::new(),
842 related_files: Some(Vec::new()),
843 active_buffer_diagnostics: Vec::new(),
844 excerpt_ranges: zeta_prompt::ExcerptRanges {
845 editable_150: 0..32,
846 editable_180: 0..32,
847 editable_350: 0..32,
848 editable_512: None,
849 editable_150_context_350: 0..48,
850 editable_180_context_350: 0..48,
851 editable_350_context_150: 20..50,
852 editable_350_context_512: None,
853 editable_350_context_1024: None,
854 context_4096: None,
855 context_8192: Some(30..excerpt.len()),
856 },
857 syntax_ranges: None,
858 in_open_source_repo: false,
859 can_collect_data: false,
860 repo_url: None,
861 };
862
863 let (stored_editable_range, stored_context_range) = zeta_prompt::excerpt_range_for_format(
864 ZetaFormat::V0327SingleFile,
865 &prompt_inputs.excerpt_ranges,
866 );
867 assert!(stored_context_range.start > stored_editable_range.start);
868
869 let (editable_range, context_range) =
870 resolved_excerpt_ranges_for_format(&prompt_inputs, ZetaFormat::V0327SingleFile);
871 assert_eq!(context_range, 0..excerpt.len());
872 assert!(editable_range.start < cursor_offset);
873 assert!(editable_range.end > cursor_offset);
874
875 let prompt = TeacherPrompt::format_prompt(
876 &Example {
877 spec: edit_prediction::example_spec::ExampleSpec {
878 name: "test".to_string(),
879 repository_url: "https://github.com/zed-industries/zed.git".to_string(),
880 revision: "HEAD".to_string(),
881 tags: Vec::new(),
882 reasoning: None,
883 uncommitted_diff: String::new(),
884 cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")),
885 cursor_position: "0:0".to_string(),
886 edit_history: String::new(),
887 expected_patches: Vec::new(),
888 rejected_patch: None,
889 telemetry: None,
890 human_feedback: Vec::new(),
891 rating: None,
892 },
893 prompt_inputs: Some(prompt_inputs),
894 prompt: None,
895 predictions: Vec::new(),
896 score: Vec::new(),
897 qa: Vec::new(),
898 zed_version: None,
899 state: None,
900 },
901 editable_range,
902 context_range,
903 );
904
905 assert!(prompt.contains(TeacherPrompt::EDITABLE_REGION_START));
906 assert!(prompt.contains(TeacherPrompt::USER_CURSOR_MARKER));
907 assert!(prompt.contains("line40"));
908 }
909}