1//! Zeta2 prompt planning and generation code shared with cloud.
2pub mod retrieval_prompt;
3
4use anyhow::{Context as _, Result, anyhow};
5use cloud_llm_client::predict_edits_v3::{
6 self, DiffPathFmt, Excerpt, Line, Point, PromptFormat, ReferencedDeclaration,
7};
8use indoc::indoc;
9use ordered_float::OrderedFloat;
10use rustc_hash::{FxHashMap, FxHashSet};
11use serde::Serialize;
12use std::cmp;
13use std::fmt::Write;
14use std::sync::Arc;
15use std::{cmp::Reverse, collections::BinaryHeap, ops::Range, path::Path};
16use strum::{EnumIter, IntoEnumIterator};
17
18pub const DEFAULT_MAX_PROMPT_BYTES: usize = 10 * 1024;
19
20pub const CURSOR_MARKER: &str = "<|user_cursor|>";
21/// NOTE: Differs from zed version of constant - includes a newline
22pub const EDITABLE_REGION_START_MARKER_WITH_NEWLINE: &str = "<|editable_region_start|>\n";
23/// NOTE: Differs from zed version of constant - includes a newline
24pub const EDITABLE_REGION_END_MARKER_WITH_NEWLINE: &str = "<|editable_region_end|>\n";
25
26// TODO: use constants for markers?
27const MARKED_EXCERPT_INSTRUCTIONS: &str = indoc! {"
28 You are a code completion assistant and your task is to analyze user edits and then rewrite an excerpt that the user provides, suggesting the appropriate edits within the excerpt, taking into account the cursor location.
29
30 The excerpt to edit will be wrapped in markers <|editable_region_start|> and <|editable_region_end|>. The cursor position is marked with <|user_cursor|>. Please respond with edited code for that region.
31
32 Other code is provided for context, and `…` indicates when code has been skipped.
33
34 # Edit History:
35
36"};
37
38const LABELED_SECTIONS_INSTRUCTIONS: &str = indoc! {r#"
39 You are a code completion assistant and your task is to analyze user edits, and suggest an edit to one of the provided sections of code.
40
41 Sections of code are grouped by file and then labeled by `<|section_N|>` (e.g `<|section_8|>`).
42
43 The cursor position is marked with `<|user_cursor|>` and it will appear within a special section labeled `<|current_section|>`. Prefer editing the current section until no more changes are needed within it.
44
45 Respond ONLY with the name of the section to edit on a single line, followed by all of the code that should replace that section. For example:
46
47 <|current_section|>
48 for i in 0..16 {
49 println!("{i}");
50 }
51
52 # Edit History:
53
54"#};
55
56const NUMBERED_LINES_INSTRUCTIONS: &str = indoc! {r#"
57 # Instructions
58
59 You are an edit prediction agent in a code editor.
60 Your job is to predict the next edit that the user will make,
61 based on their last few edits and their current cursor location.
62
63 ## Output Format
64
65 You must briefly explain your understanding of the user's goal, in one
66 or two sentences, and then specify their next edit in the form of a
67 unified diff, like this:
68
69 ```
70 --- a/src/myapp/cli.py
71 +++ b/src/myapp/cli.py
72 @@ ... @@
73 import os
74 import time
75 import sys
76 +from constants import LOG_LEVEL_WARNING
77 @@ ... @@
78 config.headless()
79 config.set_interactive(false)
80 -config.set_log_level(LOG_L)
81 +config.set_log_level(LOG_LEVEL_WARNING)
82 config.set_use_color(True)
83 ```
84
85 ## Edit History
86
87"#};
88
89const UNIFIED_DIFF_REMINDER: &str = indoc! {"
90 ---
91
92 Analyze the edit history and the files, then provide the unified diff for your predicted edits.
93 Do not include the cursor marker in your output.
94 Your diff should include edited file paths in its file headers (lines beginning with `---` and `+++`).
95 Do not include line numbers in the hunk headers, use `@@ ... @@`.
96 Removed lines begin with `-`.
97 Added lines begin with `+`.
98 Context lines begin with an extra space.
99 Context and removed lines are used to match the target edit location, so make sure to include enough of them
100 to uniquely identify it amongst all excerpts of code provided.
101"};
102
103const XML_TAGS_INSTRUCTIONS: &str = indoc! {r#"
104 # Instructions
105
106 You are an edit prediction agent in a code editor.
107 Your job is to predict the next edit that the user will make,
108 based on their last few edits and their current cursor location.
109
110 # Output Format
111
112 You must briefly explain your understanding of the user's goal, in one
113 or two sentences, and then specify their next edit, using the following
114 XML format:
115
116 <edits path="my-project/src/myapp/cli.py">
117 <old_text>
118 OLD TEXT 1 HERE
119 </old_text>
120 <new_text>
121 NEW TEXT 1 HERE
122 </new_text>
123
124 <old_text>
125 OLD TEXT 1 HERE
126 </old_text>
127 <new_text>
128 NEW TEXT 1 HERE
129 </new_text>
130 </edits>
131
132 - Specify the file to edit using the `path` attribute.
133 - Use `<old_text>` and `<new_text>` tags to replace content
134 - `<old_text>` must exactly match existing file content, including indentation
135 - `<old_text>` cannot be empty
136 - Do not escape quotes, newlines, or other characters within tags
137 - Always close all tags properly
138 - Don't include the <|user_cursor|> marker in your output.
139
140 # Edit History:
141
142"#};
143
144const OLD_TEXT_NEW_TEXT_REMINDER: &str = indoc! {r#"
145 ---
146
147 Remember that the edits in the edit history have already been deployed.
148 The files are currently as shown in the Code Excerpts section.
149"#};
150
151pub fn build_prompt(
152 request: &predict_edits_v3::PredictEditsRequest,
153) -> Result<(String, SectionLabels)> {
154 let mut insertions = match request.prompt_format {
155 PromptFormat::MarkedExcerpt => vec![
156 (
157 Point {
158 line: request.excerpt_line_range.start,
159 column: 0,
160 },
161 EDITABLE_REGION_START_MARKER_WITH_NEWLINE,
162 ),
163 (request.cursor_point, CURSOR_MARKER),
164 (
165 Point {
166 line: request.excerpt_line_range.end,
167 column: 0,
168 },
169 EDITABLE_REGION_END_MARKER_WITH_NEWLINE,
170 ),
171 ],
172 PromptFormat::LabeledSections
173 | PromptFormat::NumLinesUniDiff
174 | PromptFormat::OldTextNewText => {
175 vec![(request.cursor_point, CURSOR_MARKER)]
176 }
177 PromptFormat::OnlySnippets => vec![],
178 };
179
180 let mut prompt = match request.prompt_format {
181 PromptFormat::MarkedExcerpt => MARKED_EXCERPT_INSTRUCTIONS.to_string(),
182 PromptFormat::LabeledSections => LABELED_SECTIONS_INSTRUCTIONS.to_string(),
183 PromptFormat::NumLinesUniDiff => NUMBERED_LINES_INSTRUCTIONS.to_string(),
184 PromptFormat::OldTextNewText => XML_TAGS_INSTRUCTIONS.to_string(),
185 PromptFormat::OnlySnippets => String::new(),
186 };
187
188 if request.events.is_empty() {
189 prompt.push_str("(No edit history)\n\n");
190 } else {
191 prompt.push_str("Here are the latest edits made by the user, from earlier to later.\n\n");
192 push_events(&mut prompt, &request.events);
193 }
194
195 prompt.push_str(indoc! {"
196 # Code Excerpts
197
198 The cursor marker <|user_cursor|> indicates the current user cursor position.
199 The file is in current state, edits from edit history have been applied.
200 "});
201
202 if request.prompt_format == PromptFormat::NumLinesUniDiff {
203 prompt.push_str(indoc! {"
204 We prepend line numbers (e.g., `123|<actual line>`); they are not part of the file.
205 "});
206 }
207
208 prompt.push('\n');
209
210 let mut section_labels = Default::default();
211
212 if !request.referenced_declarations.is_empty() || !request.signatures.is_empty() {
213 let syntax_based_prompt = SyntaxBasedPrompt::populate(request)?;
214 section_labels = syntax_based_prompt.write(&mut insertions, &mut prompt)?;
215 } else {
216 if request.prompt_format == PromptFormat::LabeledSections {
217 anyhow::bail!("PromptFormat::LabeledSections cannot be used with ContextMode::Llm");
218 }
219
220 for related_file in &request.included_files {
221 write_codeblock(
222 &related_file.path,
223 &related_file.excerpts,
224 if related_file.path == request.excerpt_path {
225 &insertions
226 } else {
227 &[]
228 },
229 related_file.max_row,
230 request.prompt_format == PromptFormat::NumLinesUniDiff,
231 &mut prompt,
232 );
233 }
234 }
235
236 match request.prompt_format {
237 PromptFormat::NumLinesUniDiff => {
238 prompt.push_str(UNIFIED_DIFF_REMINDER);
239 }
240 PromptFormat::OldTextNewText => {
241 prompt.push_str(OLD_TEXT_NEW_TEXT_REMINDER);
242 }
243 _ => {}
244 }
245
246 Ok((prompt, section_labels))
247}
248
249pub fn write_codeblock<'a>(
250 path: &Path,
251 excerpts: impl IntoIterator<Item = &'a Excerpt>,
252 sorted_insertions: &[(Point, &str)],
253 file_line_count: Line,
254 include_line_numbers: bool,
255 output: &'a mut String,
256) {
257 writeln!(output, "`````{}", DiffPathFmt(path)).unwrap();
258 write_excerpts(
259 excerpts,
260 sorted_insertions,
261 file_line_count,
262 include_line_numbers,
263 output,
264 );
265 write!(output, "`````\n\n").unwrap();
266}
267
268pub fn write_excerpts<'a>(
269 excerpts: impl IntoIterator<Item = &'a Excerpt>,
270 sorted_insertions: &[(Point, &str)],
271 file_line_count: Line,
272 include_line_numbers: bool,
273 output: &mut String,
274) {
275 let mut current_row = Line(0);
276 let mut sorted_insertions = sorted_insertions.iter().peekable();
277
278 for excerpt in excerpts {
279 if excerpt.start_line > current_row {
280 writeln!(output, "…").unwrap();
281 }
282 if excerpt.text.is_empty() {
283 return;
284 }
285
286 current_row = excerpt.start_line;
287
288 for mut line in excerpt.text.lines() {
289 if include_line_numbers {
290 write!(output, "{}|", current_row.0 + 1).unwrap();
291 }
292
293 while let Some((insertion_location, insertion_marker)) = sorted_insertions.peek() {
294 match current_row.cmp(&insertion_location.line) {
295 cmp::Ordering::Equal => {
296 let (prefix, suffix) = line.split_at(insertion_location.column as usize);
297 output.push_str(prefix);
298 output.push_str(insertion_marker);
299 line = suffix;
300 sorted_insertions.next();
301 }
302 cmp::Ordering::Less => break,
303 cmp::Ordering::Greater => {
304 sorted_insertions.next();
305 break;
306 }
307 }
308 }
309 output.push_str(line);
310 output.push('\n');
311 current_row.0 += 1;
312 }
313 }
314
315 if current_row < file_line_count {
316 writeln!(output, "…").unwrap();
317 }
318}
319
320pub fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
321 if events.is_empty() {
322 return;
323 };
324
325 writeln!(output, "`````diff").unwrap();
326 for event in events {
327 writeln!(output, "{}", event).unwrap();
328 }
329 writeln!(output, "`````\n").unwrap();
330}
331
332pub struct SyntaxBasedPrompt<'a> {
333 request: &'a predict_edits_v3::PredictEditsRequest,
334 /// Snippets to include in the prompt. These may overlap - they are merged / deduplicated in
335 /// `to_prompt_string`.
336 snippets: Vec<PlannedSnippet<'a>>,
337 budget_used: usize,
338}
339
340#[derive(Clone, Debug)]
341pub struct PlannedSnippet<'a> {
342 path: Arc<Path>,
343 range: Range<Line>,
344 text: &'a str,
345 // TODO: Indicate this in the output
346 #[allow(dead_code)]
347 text_is_truncated: bool,
348}
349
350#[derive(EnumIter, Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
351pub enum DeclarationStyle {
352 Signature,
353 Declaration,
354}
355
356#[derive(Default, Clone, Debug, Serialize)]
357pub struct SectionLabels {
358 pub excerpt_index: usize,
359 pub section_ranges: Vec<(Arc<Path>, Range<Line>)>,
360}
361
362impl<'a> SyntaxBasedPrompt<'a> {
363 /// Greedy one-pass knapsack algorithm to populate the prompt plan. Does the following:
364 ///
365 /// Initializes a priority queue by populating it with each snippet, finding the
366 /// DeclarationStyle that minimizes `score_density = score / snippet.range(style).len()`. When a
367 /// "signature" snippet is popped, insert an entry for the "declaration" variant that reflects
368 /// the cost of upgrade.
369 ///
370 /// TODO: Implement an early halting condition. One option might be to have another priority
371 /// queue where the score is the size, and update it accordingly. Another option might be to
372 /// have some simpler heuristic like bailing after N failed insertions, or based on how much
373 /// budget is left.
374 ///
375 /// TODO: Has the current known sources of imprecision:
376 ///
377 /// * Does not consider snippet overlap when ranking. For example, it might add a field to the
378 /// plan even though the containing struct is already included.
379 ///
380 /// * Does not consider cost of signatures when ranking snippets - this is tricky since
381 /// signatures may be shared by multiple snippets.
382 ///
383 /// * Does not include file paths / other text when considering max_bytes.
384 pub fn populate(request: &'a predict_edits_v3::PredictEditsRequest) -> Result<Self> {
385 let mut this = Self {
386 request,
387 snippets: Vec::new(),
388 budget_used: request.excerpt.len(),
389 };
390 let mut included_parents = FxHashSet::default();
391 let additional_parents = this.additional_parent_signatures(
392 &request.excerpt_path,
393 request.excerpt_parent,
394 &included_parents,
395 )?;
396 this.add_parents(&mut included_parents, additional_parents);
397
398 let max_bytes = request.prompt_max_bytes.unwrap_or(DEFAULT_MAX_PROMPT_BYTES);
399
400 if this.budget_used > max_bytes {
401 return Err(anyhow!(
402 "Excerpt + signatures size of {} already exceeds budget of {}",
403 this.budget_used,
404 max_bytes
405 ));
406 }
407
408 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
409 struct QueueEntry {
410 score_density: OrderedFloat<f32>,
411 declaration_index: usize,
412 style: DeclarationStyle,
413 }
414
415 // Initialize priority queue with the best score for each snippet.
416 let mut queue: BinaryHeap<QueueEntry> = BinaryHeap::new();
417 for (declaration_index, declaration) in request.referenced_declarations.iter().enumerate() {
418 let (style, score_density) = DeclarationStyle::iter()
419 .map(|style| {
420 (
421 style,
422 OrderedFloat(declaration_score_density(&declaration, style)),
423 )
424 })
425 .max_by_key(|(_, score_density)| *score_density)
426 .unwrap();
427 queue.push(QueueEntry {
428 score_density,
429 declaration_index,
430 style,
431 });
432 }
433
434 // Knapsack selection loop
435 while let Some(queue_entry) = queue.pop() {
436 let Some(declaration) = request
437 .referenced_declarations
438 .get(queue_entry.declaration_index)
439 else {
440 return Err(anyhow!(
441 "Invalid declaration index {}",
442 queue_entry.declaration_index
443 ));
444 };
445
446 let mut additional_bytes = declaration_size(declaration, queue_entry.style);
447 if this.budget_used + additional_bytes > max_bytes {
448 continue;
449 }
450
451 let additional_parents = this.additional_parent_signatures(
452 &declaration.path,
453 declaration.parent_index,
454 &mut included_parents,
455 )?;
456 additional_bytes += additional_parents
457 .iter()
458 .map(|(_, snippet)| snippet.text.len())
459 .sum::<usize>();
460 if this.budget_used + additional_bytes > max_bytes {
461 continue;
462 }
463
464 this.budget_used += additional_bytes;
465 this.add_parents(&mut included_parents, additional_parents);
466 let planned_snippet = match queue_entry.style {
467 DeclarationStyle::Signature => {
468 let Some(text) = declaration.text.get(declaration.signature_range.clone())
469 else {
470 return Err(anyhow!(
471 "Invalid declaration signature_range {:?} with text.len() = {}",
472 declaration.signature_range,
473 declaration.text.len()
474 ));
475 };
476 let signature_start_line = declaration.range.start
477 + Line(
478 declaration.text[..declaration.signature_range.start]
479 .lines()
480 .count() as u32,
481 );
482 let signature_end_line = signature_start_line
483 + Line(
484 declaration.text
485 [declaration.signature_range.start..declaration.signature_range.end]
486 .lines()
487 .count() as u32,
488 );
489 let range = signature_start_line..signature_end_line;
490
491 PlannedSnippet {
492 path: declaration.path.clone(),
493 range,
494 text,
495 text_is_truncated: declaration.text_is_truncated,
496 }
497 }
498 DeclarationStyle::Declaration => PlannedSnippet {
499 path: declaration.path.clone(),
500 range: declaration.range.clone(),
501 text: &declaration.text,
502 text_is_truncated: declaration.text_is_truncated,
503 },
504 };
505 this.snippets.push(planned_snippet);
506
507 // When a Signature is consumed, insert an entry for Definition style.
508 if queue_entry.style == DeclarationStyle::Signature {
509 let signature_size = declaration_size(&declaration, DeclarationStyle::Signature);
510 let declaration_size =
511 declaration_size(&declaration, DeclarationStyle::Declaration);
512 let signature_score = declaration_score(&declaration, DeclarationStyle::Signature);
513 let declaration_score =
514 declaration_score(&declaration, DeclarationStyle::Declaration);
515
516 let score_diff = declaration_score - signature_score;
517 let size_diff = declaration_size.saturating_sub(signature_size);
518 if score_diff > 0.0001 && size_diff > 0 {
519 queue.push(QueueEntry {
520 declaration_index: queue_entry.declaration_index,
521 score_density: OrderedFloat(score_diff / (size_diff as f32)),
522 style: DeclarationStyle::Declaration,
523 });
524 }
525 }
526 }
527
528 anyhow::Ok(this)
529 }
530
531 fn add_parents(
532 &mut self,
533 included_parents: &mut FxHashSet<usize>,
534 snippets: Vec<(usize, PlannedSnippet<'a>)>,
535 ) {
536 for (parent_index, snippet) in snippets {
537 included_parents.insert(parent_index);
538 self.budget_used += snippet.text.len();
539 self.snippets.push(snippet);
540 }
541 }
542
543 fn additional_parent_signatures(
544 &self,
545 path: &Arc<Path>,
546 parent_index: Option<usize>,
547 included_parents: &FxHashSet<usize>,
548 ) -> Result<Vec<(usize, PlannedSnippet<'a>)>> {
549 let mut results = Vec::new();
550 self.additional_parent_signatures_impl(path, parent_index, included_parents, &mut results)?;
551 Ok(results)
552 }
553
554 fn additional_parent_signatures_impl(
555 &self,
556 path: &Arc<Path>,
557 parent_index: Option<usize>,
558 included_parents: &FxHashSet<usize>,
559 results: &mut Vec<(usize, PlannedSnippet<'a>)>,
560 ) -> Result<()> {
561 let Some(parent_index) = parent_index else {
562 return Ok(());
563 };
564 if included_parents.contains(&parent_index) {
565 return Ok(());
566 }
567 let Some(parent_signature) = self.request.signatures.get(parent_index) else {
568 return Err(anyhow!("Invalid parent index {}", parent_index));
569 };
570 results.push((
571 parent_index,
572 PlannedSnippet {
573 path: path.clone(),
574 range: parent_signature.range.clone(),
575 text: &parent_signature.text,
576 text_is_truncated: parent_signature.text_is_truncated,
577 },
578 ));
579 self.additional_parent_signatures_impl(
580 path,
581 parent_signature.parent_index,
582 included_parents,
583 results,
584 )
585 }
586
587 /// Renders the planned context. Each file starts with "```FILE_PATH\n` and ends with triple
588 /// backticks, with a newline after each file. Outputs a line with "..." between nonconsecutive
589 /// chunks.
590 pub fn write(
591 &'a self,
592 excerpt_file_insertions: &mut Vec<(Point, &'static str)>,
593 prompt: &mut String,
594 ) -> Result<SectionLabels> {
595 let mut file_to_snippets: FxHashMap<&'a std::path::Path, Vec<&PlannedSnippet<'a>>> =
596 FxHashMap::default();
597 for snippet in &self.snippets {
598 file_to_snippets
599 .entry(&snippet.path)
600 .or_default()
601 .push(snippet);
602 }
603
604 // Reorder so that file with cursor comes last
605 let mut file_snippets = Vec::new();
606 let mut excerpt_file_snippets = Vec::new();
607 for (file_path, snippets) in file_to_snippets {
608 if file_path == self.request.excerpt_path.as_ref() {
609 excerpt_file_snippets = snippets;
610 } else {
611 file_snippets.push((file_path, snippets, false));
612 }
613 }
614 let excerpt_snippet = PlannedSnippet {
615 path: self.request.excerpt_path.clone(),
616 range: self.request.excerpt_line_range.clone(),
617 text: &self.request.excerpt,
618 text_is_truncated: false,
619 };
620 excerpt_file_snippets.push(&excerpt_snippet);
621 file_snippets.push((&self.request.excerpt_path, excerpt_file_snippets, true));
622
623 let section_labels =
624 self.push_file_snippets(prompt, excerpt_file_insertions, file_snippets)?;
625
626 Ok(section_labels)
627 }
628
629 fn push_file_snippets(
630 &self,
631 output: &mut String,
632 excerpt_file_insertions: &mut Vec<(Point, &'static str)>,
633 file_snippets: Vec<(&'a Path, Vec<&'a PlannedSnippet>, bool)>,
634 ) -> Result<SectionLabels> {
635 let mut section_ranges = Vec::new();
636 let mut excerpt_index = None;
637
638 for (file_path, mut snippets, is_excerpt_file) in file_snippets {
639 snippets.sort_by_key(|s| (s.range.start, Reverse(s.range.end)));
640
641 // TODO: What if the snippets get expanded too large to be editable?
642 let mut current_snippet: Option<(&PlannedSnippet, Range<Line>)> = None;
643 let mut disjoint_snippets: Vec<(&PlannedSnippet, Range<Line>)> = Vec::new();
644 for snippet in snippets {
645 if let Some((_, current_snippet_range)) = current_snippet.as_mut()
646 && snippet.range.start <= current_snippet_range.end
647 {
648 current_snippet_range.end = current_snippet_range.end.max(snippet.range.end);
649 continue;
650 }
651 if let Some(current_snippet) = current_snippet.take() {
652 disjoint_snippets.push(current_snippet);
653 }
654 current_snippet = Some((snippet, snippet.range.clone()));
655 }
656 if let Some(current_snippet) = current_snippet.take() {
657 disjoint_snippets.push(current_snippet);
658 }
659
660 writeln!(output, "`````path={}", file_path.display()).ok();
661 let mut skipped_last_snippet = false;
662 for (snippet, range) in disjoint_snippets {
663 let section_index = section_ranges.len();
664
665 match self.request.prompt_format {
666 PromptFormat::MarkedExcerpt
667 | PromptFormat::OnlySnippets
668 | PromptFormat::OldTextNewText
669 | PromptFormat::NumLinesUniDiff => {
670 if range.start.0 > 0 && !skipped_last_snippet {
671 output.push_str("…\n");
672 }
673 }
674 PromptFormat::LabeledSections => {
675 if is_excerpt_file
676 && range.start <= self.request.excerpt_line_range.start
677 && range.end >= self.request.excerpt_line_range.end
678 {
679 writeln!(output, "<|current_section|>").ok();
680 } else {
681 writeln!(output, "<|section_{}|>", section_index).ok();
682 }
683 }
684 }
685
686 let push_full_snippet = |output: &mut String| {
687 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
688 for (i, line) in snippet.text.lines().enumerate() {
689 writeln!(output, "{}|{}", i as u32 + range.start.0 + 1, line)?;
690 }
691 } else {
692 output.push_str(&snippet.text);
693 }
694 anyhow::Ok(())
695 };
696
697 if is_excerpt_file {
698 if self.request.prompt_format == PromptFormat::OnlySnippets {
699 if range.start >= self.request.excerpt_line_range.start
700 && range.end <= self.request.excerpt_line_range.end
701 {
702 skipped_last_snippet = true;
703 } else {
704 skipped_last_snippet = false;
705 output.push_str(snippet.text);
706 }
707 } else if !excerpt_file_insertions.is_empty() {
708 let lines = snippet.text.lines().collect::<Vec<_>>();
709 let push_line = |output: &mut String, line_ix: usize| {
710 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
711 write!(output, "{}|", line_ix as u32 + range.start.0 + 1)?;
712 }
713 anyhow::Ok(writeln!(output, "{}", lines[line_ix])?)
714 };
715 let mut last_line_ix = 0;
716 let mut insertion_ix = 0;
717 while insertion_ix < excerpt_file_insertions.len() {
718 let (point, insertion) = &excerpt_file_insertions[insertion_ix];
719 let found = point.line >= range.start && point.line <= range.end;
720 if found {
721 excerpt_index = Some(section_index);
722 let insertion_line_ix = (point.line.0 - range.start.0) as usize;
723 for line_ix in last_line_ix..insertion_line_ix {
724 push_line(output, line_ix)?;
725 }
726 if let Some(next_line) = lines.get(insertion_line_ix) {
727 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
728 write!(
729 output,
730 "{}|",
731 insertion_line_ix as u32 + range.start.0 + 1
732 )?
733 }
734 output.push_str(&next_line[..point.column as usize]);
735 output.push_str(insertion);
736 writeln!(output, "{}", &next_line[point.column as usize..])?;
737 } else {
738 writeln!(output, "{}", insertion)?;
739 }
740 last_line_ix = insertion_line_ix + 1;
741 excerpt_file_insertions.remove(insertion_ix);
742 continue;
743 }
744 insertion_ix += 1;
745 }
746 skipped_last_snippet = false;
747 for line_ix in last_line_ix..lines.len() {
748 push_line(output, line_ix)?;
749 }
750 } else {
751 skipped_last_snippet = false;
752 push_full_snippet(output)?;
753 }
754 } else {
755 skipped_last_snippet = false;
756 push_full_snippet(output)?;
757 }
758
759 section_ranges.push((snippet.path.clone(), range));
760 }
761
762 output.push_str("`````\n\n");
763 }
764
765 Ok(SectionLabels {
766 // TODO: Clean this up
767 excerpt_index: match self.request.prompt_format {
768 PromptFormat::OnlySnippets => 0,
769 _ => excerpt_index.context("bug: no snippet found for excerpt")?,
770 },
771 section_ranges,
772 })
773 }
774}
775
776fn declaration_score_density(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> f32 {
777 declaration_score(declaration, style) / declaration_size(declaration, style) as f32
778}
779
780fn declaration_score(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> f32 {
781 match style {
782 DeclarationStyle::Signature => declaration.signature_score,
783 DeclarationStyle::Declaration => declaration.declaration_score,
784 }
785}
786
787fn declaration_size(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> usize {
788 match style {
789 DeclarationStyle::Signature => declaration.signature_range.len(),
790 DeclarationStyle::Declaration => declaration.text.len(),
791 }
792}