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