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 a code completion assistant helping a programmer finish their work. Your task is to:
60
61 1. Analyze the edit history to understand what the programmer is trying to achieve
62 2. Identify any incomplete refactoring or changes that need to be finished
63 3. Make the remaining edits that a human programmer would logically make next
64 4. Apply systematic changes consistently across the entire codebase - if you see a pattern starting, complete it everywhere.
65
66 Focus on:
67 - Understanding the intent behind the changes (e.g., improving error handling, refactoring APIs, fixing bugs)
68 - Completing any partially-applied changes across the codebase
69 - Ensuring consistency with the programming style and patterns already established
70 - Making edits that maintain or improve code quality
71 - If the programmer started refactoring one instance of a pattern, find and update ALL similar instances
72 - Don't write a lot of code if you're not sure what to do
73
74 Rules:
75 - Do not just mechanically apply patterns - reason about what changes make sense given the context and the programmer's apparent goals.
76 - Do not just fix syntax errors - look for the broader refactoring pattern and apply it systematically throughout the code.
77 - Write the edits in the unified diff format as shown in the example.
78
79 # Example output:
80
81 ```
82 --- a/src/myapp/cli.py
83 +++ b/src/myapp/cli.py
84 @@ -1,3 +1,3 @@
85 -
86 -
87 -import sys
88 +import json
89 ```
90
91 # Edit History:
92
93"#};
94
95const UNIFIED_DIFF_REMINDER: &str = indoc! {"
96 ---
97
98 Please analyze the edit history and the files, then provide the unified diff for your predicted edits.
99 Do not include the cursor marker in your output.
100 If you're editing multiple files, be sure to reflect filename in the hunk's header.
101"};
102
103pub fn build_prompt(
104 request: &predict_edits_v3::PredictEditsRequest,
105) -> Result<(String, SectionLabels)> {
106 let mut insertions = match request.prompt_format {
107 PromptFormat::MarkedExcerpt => vec![
108 (
109 Point {
110 line: request.excerpt_line_range.start,
111 column: 0,
112 },
113 EDITABLE_REGION_START_MARKER_WITH_NEWLINE,
114 ),
115 (request.cursor_point, CURSOR_MARKER),
116 (
117 Point {
118 line: request.excerpt_line_range.end,
119 column: 0,
120 },
121 EDITABLE_REGION_END_MARKER_WITH_NEWLINE,
122 ),
123 ],
124 PromptFormat::LabeledSections => vec![(request.cursor_point, CURSOR_MARKER)],
125 PromptFormat::NumLinesUniDiff => {
126 vec![(request.cursor_point, CURSOR_MARKER)]
127 }
128 PromptFormat::OnlySnippets => vec![],
129 };
130
131 let mut prompt = match request.prompt_format {
132 PromptFormat::MarkedExcerpt => MARKED_EXCERPT_INSTRUCTIONS.to_string(),
133 PromptFormat::LabeledSections => LABELED_SECTIONS_INSTRUCTIONS.to_string(),
134 PromptFormat::NumLinesUniDiff => NUMBERED_LINES_INSTRUCTIONS.to_string(),
135 // only intended for use via zeta_cli
136 PromptFormat::OnlySnippets => String::new(),
137 };
138
139 if request.events.is_empty() {
140 prompt.push_str("(No edit history)\n\n");
141 } else {
142 prompt.push_str(
143 "The following are the latest edits made by the user, from earlier to later.\n\n",
144 );
145 push_events(&mut prompt, &request.events);
146 }
147
148 if request.prompt_format == PromptFormat::NumLinesUniDiff {
149 if request.referenced_declarations.is_empty() {
150 prompt.push_str(indoc! {"
151 # File under the cursor:
152
153 The cursor marker <|user_cursor|> indicates the current user cursor position.
154 The file is in current state, edits from edit history have been applied.
155 We prepend line numbers (e.g., `123|<actual line>`); they are not part of the file.
156
157 "});
158 } else {
159 // Note: This hasn't been trained on yet
160 prompt.push_str(indoc! {"
161 # Code Excerpts:
162
163 The cursor marker <|user_cursor|> indicates the current user cursor position.
164 Other excerpts of code from the project have been included as context based on their similarity to the code under the cursor.
165 Context excerpts are not guaranteed to be relevant, so use your own judgement.
166 Files are in their current state, edits from edit history have been applied.
167 We prepend line numbers (e.g., `123|<actual line>`); they are not part of the file.
168
169 "});
170 }
171 } else {
172 prompt.push_str("\n## Code\n\n");
173 }
174
175 let mut section_labels = Default::default();
176
177 if !request.referenced_declarations.is_empty() || !request.signatures.is_empty() {
178 let syntax_based_prompt = SyntaxBasedPrompt::populate(request)?;
179 section_labels = syntax_based_prompt.write(&mut insertions, &mut prompt)?;
180 } else {
181 if request.prompt_format == PromptFormat::LabeledSections {
182 anyhow::bail!("PromptFormat::LabeledSections cannot be used with ContextMode::Llm");
183 }
184
185 for related_file in &request.included_files {
186 write_codeblock(
187 &related_file.path,
188 &related_file.excerpts,
189 if related_file.path == request.excerpt_path {
190 &insertions
191 } else {
192 &[]
193 },
194 related_file.max_row,
195 request.prompt_format == PromptFormat::NumLinesUniDiff,
196 &mut prompt,
197 );
198 }
199 }
200
201 if request.prompt_format == PromptFormat::NumLinesUniDiff {
202 prompt.push_str(UNIFIED_DIFF_REMINDER);
203 }
204
205 Ok((prompt, section_labels))
206}
207
208pub fn write_codeblock<'a>(
209 path: &Path,
210 excerpts: impl IntoIterator<Item = &'a Excerpt>,
211 sorted_insertions: &[(Point, &str)],
212 file_line_count: Line,
213 include_line_numbers: bool,
214 output: &'a mut String,
215) {
216 writeln!(output, "`````{}", DiffPathFmt(path)).unwrap();
217 write_excerpts(
218 excerpts,
219 sorted_insertions,
220 file_line_count,
221 include_line_numbers,
222 output,
223 );
224 write!(output, "`````\n\n").unwrap();
225}
226
227pub fn write_excerpts<'a>(
228 excerpts: impl IntoIterator<Item = &'a Excerpt>,
229 sorted_insertions: &[(Point, &str)],
230 file_line_count: Line,
231 include_line_numbers: bool,
232 output: &mut String,
233) {
234 let mut current_row = Line(0);
235 let mut sorted_insertions = sorted_insertions.iter().peekable();
236
237 for excerpt in excerpts {
238 if excerpt.start_line > current_row {
239 writeln!(output, "…").unwrap();
240 }
241 if excerpt.text.is_empty() {
242 return;
243 }
244
245 current_row = excerpt.start_line;
246
247 for mut line in excerpt.text.lines() {
248 if include_line_numbers {
249 write!(output, "{}|", current_row.0 + 1).unwrap();
250 }
251
252 while let Some((insertion_location, insertion_marker)) = sorted_insertions.peek() {
253 match current_row.cmp(&insertion_location.line) {
254 cmp::Ordering::Equal => {
255 let (prefix, suffix) = line.split_at(insertion_location.column as usize);
256 output.push_str(prefix);
257 output.push_str(insertion_marker);
258 line = suffix;
259 sorted_insertions.next();
260 }
261 cmp::Ordering::Less => break,
262 cmp::Ordering::Greater => {
263 sorted_insertions.next();
264 break;
265 }
266 }
267 }
268 output.push_str(line);
269 output.push('\n');
270 current_row.0 += 1;
271 }
272 }
273
274 if current_row < file_line_count {
275 writeln!(output, "…").unwrap();
276 }
277}
278
279pub fn push_events(output: &mut String, events: &[predict_edits_v3::Event]) {
280 if events.is_empty() {
281 return;
282 };
283
284 writeln!(output, "`````diff").unwrap();
285 for event in events {
286 writeln!(output, "{}", event).unwrap();
287 }
288 writeln!(output, "`````\n").unwrap();
289}
290
291pub struct SyntaxBasedPrompt<'a> {
292 request: &'a predict_edits_v3::PredictEditsRequest,
293 /// Snippets to include in the prompt. These may overlap - they are merged / deduplicated in
294 /// `to_prompt_string`.
295 snippets: Vec<PlannedSnippet<'a>>,
296 budget_used: usize,
297}
298
299#[derive(Clone, Debug)]
300pub struct PlannedSnippet<'a> {
301 path: Arc<Path>,
302 range: Range<Line>,
303 text: &'a str,
304 // TODO: Indicate this in the output
305 #[allow(dead_code)]
306 text_is_truncated: bool,
307}
308
309#[derive(EnumIter, Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
310pub enum DeclarationStyle {
311 Signature,
312 Declaration,
313}
314
315#[derive(Default, Clone, Debug, Serialize)]
316pub struct SectionLabels {
317 pub excerpt_index: usize,
318 pub section_ranges: Vec<(Arc<Path>, Range<Line>)>,
319}
320
321impl<'a> SyntaxBasedPrompt<'a> {
322 /// Greedy one-pass knapsack algorithm to populate the prompt plan. Does the following:
323 ///
324 /// Initializes a priority queue by populating it with each snippet, finding the
325 /// DeclarationStyle that minimizes `score_density = score / snippet.range(style).len()`. When a
326 /// "signature" snippet is popped, insert an entry for the "declaration" variant that reflects
327 /// the cost of upgrade.
328 ///
329 /// TODO: Implement an early halting condition. One option might be to have another priority
330 /// queue where the score is the size, and update it accordingly. Another option might be to
331 /// have some simpler heuristic like bailing after N failed insertions, or based on how much
332 /// budget is left.
333 ///
334 /// TODO: Has the current known sources of imprecision:
335 ///
336 /// * Does not consider snippet overlap when ranking. For example, it might add a field to the
337 /// plan even though the containing struct is already included.
338 ///
339 /// * Does not consider cost of signatures when ranking snippets - this is tricky since
340 /// signatures may be shared by multiple snippets.
341 ///
342 /// * Does not include file paths / other text when considering max_bytes.
343 pub fn populate(request: &'a predict_edits_v3::PredictEditsRequest) -> Result<Self> {
344 let mut this = Self {
345 request,
346 snippets: Vec::new(),
347 budget_used: request.excerpt.len(),
348 };
349 let mut included_parents = FxHashSet::default();
350 let additional_parents = this.additional_parent_signatures(
351 &request.excerpt_path,
352 request.excerpt_parent,
353 &included_parents,
354 )?;
355 this.add_parents(&mut included_parents, additional_parents);
356
357 let max_bytes = request.prompt_max_bytes.unwrap_or(DEFAULT_MAX_PROMPT_BYTES);
358
359 if this.budget_used > max_bytes {
360 return Err(anyhow!(
361 "Excerpt + signatures size of {} already exceeds budget of {}",
362 this.budget_used,
363 max_bytes
364 ));
365 }
366
367 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
368 struct QueueEntry {
369 score_density: OrderedFloat<f32>,
370 declaration_index: usize,
371 style: DeclarationStyle,
372 }
373
374 // Initialize priority queue with the best score for each snippet.
375 let mut queue: BinaryHeap<QueueEntry> = BinaryHeap::new();
376 for (declaration_index, declaration) in request.referenced_declarations.iter().enumerate() {
377 let (style, score_density) = DeclarationStyle::iter()
378 .map(|style| {
379 (
380 style,
381 OrderedFloat(declaration_score_density(&declaration, style)),
382 )
383 })
384 .max_by_key(|(_, score_density)| *score_density)
385 .unwrap();
386 queue.push(QueueEntry {
387 score_density,
388 declaration_index,
389 style,
390 });
391 }
392
393 // Knapsack selection loop
394 while let Some(queue_entry) = queue.pop() {
395 let Some(declaration) = request
396 .referenced_declarations
397 .get(queue_entry.declaration_index)
398 else {
399 return Err(anyhow!(
400 "Invalid declaration index {}",
401 queue_entry.declaration_index
402 ));
403 };
404
405 let mut additional_bytes = declaration_size(declaration, queue_entry.style);
406 if this.budget_used + additional_bytes > max_bytes {
407 continue;
408 }
409
410 let additional_parents = this.additional_parent_signatures(
411 &declaration.path,
412 declaration.parent_index,
413 &mut included_parents,
414 )?;
415 additional_bytes += additional_parents
416 .iter()
417 .map(|(_, snippet)| snippet.text.len())
418 .sum::<usize>();
419 if this.budget_used + additional_bytes > max_bytes {
420 continue;
421 }
422
423 this.budget_used += additional_bytes;
424 this.add_parents(&mut included_parents, additional_parents);
425 let planned_snippet = match queue_entry.style {
426 DeclarationStyle::Signature => {
427 let Some(text) = declaration.text.get(declaration.signature_range.clone())
428 else {
429 return Err(anyhow!(
430 "Invalid declaration signature_range {:?} with text.len() = {}",
431 declaration.signature_range,
432 declaration.text.len()
433 ));
434 };
435 let signature_start_line = declaration.range.start
436 + Line(
437 declaration.text[..declaration.signature_range.start]
438 .lines()
439 .count() as u32,
440 );
441 let signature_end_line = signature_start_line
442 + Line(
443 declaration.text
444 [declaration.signature_range.start..declaration.signature_range.end]
445 .lines()
446 .count() as u32,
447 );
448 let range = signature_start_line..signature_end_line;
449
450 PlannedSnippet {
451 path: declaration.path.clone(),
452 range,
453 text,
454 text_is_truncated: declaration.text_is_truncated,
455 }
456 }
457 DeclarationStyle::Declaration => PlannedSnippet {
458 path: declaration.path.clone(),
459 range: declaration.range.clone(),
460 text: &declaration.text,
461 text_is_truncated: declaration.text_is_truncated,
462 },
463 };
464 this.snippets.push(planned_snippet);
465
466 // When a Signature is consumed, insert an entry for Definition style.
467 if queue_entry.style == DeclarationStyle::Signature {
468 let signature_size = declaration_size(&declaration, DeclarationStyle::Signature);
469 let declaration_size =
470 declaration_size(&declaration, DeclarationStyle::Declaration);
471 let signature_score = declaration_score(&declaration, DeclarationStyle::Signature);
472 let declaration_score =
473 declaration_score(&declaration, DeclarationStyle::Declaration);
474
475 let score_diff = declaration_score - signature_score;
476 let size_diff = declaration_size.saturating_sub(signature_size);
477 if score_diff > 0.0001 && size_diff > 0 {
478 queue.push(QueueEntry {
479 declaration_index: queue_entry.declaration_index,
480 score_density: OrderedFloat(score_diff / (size_diff as f32)),
481 style: DeclarationStyle::Declaration,
482 });
483 }
484 }
485 }
486
487 anyhow::Ok(this)
488 }
489
490 fn add_parents(
491 &mut self,
492 included_parents: &mut FxHashSet<usize>,
493 snippets: Vec<(usize, PlannedSnippet<'a>)>,
494 ) {
495 for (parent_index, snippet) in snippets {
496 included_parents.insert(parent_index);
497 self.budget_used += snippet.text.len();
498 self.snippets.push(snippet);
499 }
500 }
501
502 fn additional_parent_signatures(
503 &self,
504 path: &Arc<Path>,
505 parent_index: Option<usize>,
506 included_parents: &FxHashSet<usize>,
507 ) -> Result<Vec<(usize, PlannedSnippet<'a>)>> {
508 let mut results = Vec::new();
509 self.additional_parent_signatures_impl(path, parent_index, included_parents, &mut results)?;
510 Ok(results)
511 }
512
513 fn additional_parent_signatures_impl(
514 &self,
515 path: &Arc<Path>,
516 parent_index: Option<usize>,
517 included_parents: &FxHashSet<usize>,
518 results: &mut Vec<(usize, PlannedSnippet<'a>)>,
519 ) -> Result<()> {
520 let Some(parent_index) = parent_index else {
521 return Ok(());
522 };
523 if included_parents.contains(&parent_index) {
524 return Ok(());
525 }
526 let Some(parent_signature) = self.request.signatures.get(parent_index) else {
527 return Err(anyhow!("Invalid parent index {}", parent_index));
528 };
529 results.push((
530 parent_index,
531 PlannedSnippet {
532 path: path.clone(),
533 range: parent_signature.range.clone(),
534 text: &parent_signature.text,
535 text_is_truncated: parent_signature.text_is_truncated,
536 },
537 ));
538 self.additional_parent_signatures_impl(
539 path,
540 parent_signature.parent_index,
541 included_parents,
542 results,
543 )
544 }
545
546 /// Renders the planned context. Each file starts with "```FILE_PATH\n` and ends with triple
547 /// backticks, with a newline after each file. Outputs a line with "..." between nonconsecutive
548 /// chunks.
549 pub fn write(
550 &'a self,
551 excerpt_file_insertions: &mut Vec<(Point, &'static str)>,
552 prompt: &mut String,
553 ) -> Result<SectionLabels> {
554 let mut file_to_snippets: FxHashMap<&'a std::path::Path, Vec<&PlannedSnippet<'a>>> =
555 FxHashMap::default();
556 for snippet in &self.snippets {
557 file_to_snippets
558 .entry(&snippet.path)
559 .or_default()
560 .push(snippet);
561 }
562
563 // Reorder so that file with cursor comes last
564 let mut file_snippets = Vec::new();
565 let mut excerpt_file_snippets = Vec::new();
566 for (file_path, snippets) in file_to_snippets {
567 if file_path == self.request.excerpt_path.as_ref() {
568 excerpt_file_snippets = snippets;
569 } else {
570 file_snippets.push((file_path, snippets, false));
571 }
572 }
573 let excerpt_snippet = PlannedSnippet {
574 path: self.request.excerpt_path.clone(),
575 range: self.request.excerpt_line_range.clone(),
576 text: &self.request.excerpt,
577 text_is_truncated: false,
578 };
579 excerpt_file_snippets.push(&excerpt_snippet);
580 file_snippets.push((&self.request.excerpt_path, excerpt_file_snippets, true));
581
582 let section_labels =
583 self.push_file_snippets(prompt, excerpt_file_insertions, file_snippets)?;
584
585 Ok(section_labels)
586 }
587
588 fn push_file_snippets(
589 &self,
590 output: &mut String,
591 excerpt_file_insertions: &mut Vec<(Point, &'static str)>,
592 file_snippets: Vec<(&'a Path, Vec<&'a PlannedSnippet>, bool)>,
593 ) -> Result<SectionLabels> {
594 let mut section_ranges = Vec::new();
595 let mut excerpt_index = None;
596
597 for (file_path, mut snippets, is_excerpt_file) in file_snippets {
598 snippets.sort_by_key(|s| (s.range.start, Reverse(s.range.end)));
599
600 // TODO: What if the snippets get expanded too large to be editable?
601 let mut current_snippet: Option<(&PlannedSnippet, Range<Line>)> = None;
602 let mut disjoint_snippets: Vec<(&PlannedSnippet, Range<Line>)> = Vec::new();
603 for snippet in snippets {
604 if let Some((_, current_snippet_range)) = current_snippet.as_mut()
605 && snippet.range.start <= current_snippet_range.end
606 {
607 current_snippet_range.end = current_snippet_range.end.max(snippet.range.end);
608 continue;
609 }
610 if let Some(current_snippet) = current_snippet.take() {
611 disjoint_snippets.push(current_snippet);
612 }
613 current_snippet = Some((snippet, snippet.range.clone()));
614 }
615 if let Some(current_snippet) = current_snippet.take() {
616 disjoint_snippets.push(current_snippet);
617 }
618
619 writeln!(output, "`````path={}", file_path.display()).ok();
620 let mut skipped_last_snippet = false;
621 for (snippet, range) in disjoint_snippets {
622 let section_index = section_ranges.len();
623
624 match self.request.prompt_format {
625 PromptFormat::MarkedExcerpt
626 | PromptFormat::OnlySnippets
627 | PromptFormat::NumLinesUniDiff => {
628 if range.start.0 > 0 && !skipped_last_snippet {
629 output.push_str("…\n");
630 }
631 }
632 PromptFormat::LabeledSections => {
633 if is_excerpt_file
634 && range.start <= self.request.excerpt_line_range.start
635 && range.end >= self.request.excerpt_line_range.end
636 {
637 writeln!(output, "<|current_section|>").ok();
638 } else {
639 writeln!(output, "<|section_{}|>", section_index).ok();
640 }
641 }
642 }
643
644 let push_full_snippet = |output: &mut String| {
645 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
646 for (i, line) in snippet.text.lines().enumerate() {
647 writeln!(output, "{}|{}", i as u32 + range.start.0 + 1, line)?;
648 }
649 } else {
650 output.push_str(&snippet.text);
651 }
652 anyhow::Ok(())
653 };
654
655 if is_excerpt_file {
656 if self.request.prompt_format == PromptFormat::OnlySnippets {
657 if range.start >= self.request.excerpt_line_range.start
658 && range.end <= self.request.excerpt_line_range.end
659 {
660 skipped_last_snippet = true;
661 } else {
662 skipped_last_snippet = false;
663 output.push_str(snippet.text);
664 }
665 } else if !excerpt_file_insertions.is_empty() {
666 let lines = snippet.text.lines().collect::<Vec<_>>();
667 let push_line = |output: &mut String, line_ix: usize| {
668 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
669 write!(output, "{}|", line_ix as u32 + range.start.0 + 1)?;
670 }
671 anyhow::Ok(writeln!(output, "{}", lines[line_ix])?)
672 };
673 let mut last_line_ix = 0;
674 let mut insertion_ix = 0;
675 while insertion_ix < excerpt_file_insertions.len() {
676 let (point, insertion) = &excerpt_file_insertions[insertion_ix];
677 let found = point.line >= range.start && point.line <= range.end;
678 if found {
679 excerpt_index = Some(section_index);
680 let insertion_line_ix = (point.line.0 - range.start.0) as usize;
681 for line_ix in last_line_ix..insertion_line_ix {
682 push_line(output, line_ix)?;
683 }
684 if let Some(next_line) = lines.get(insertion_line_ix) {
685 if self.request.prompt_format == PromptFormat::NumLinesUniDiff {
686 write!(
687 output,
688 "{}|",
689 insertion_line_ix as u32 + range.start.0 + 1
690 )?
691 }
692 output.push_str(&next_line[..point.column as usize]);
693 output.push_str(insertion);
694 writeln!(output, "{}", &next_line[point.column as usize..])?;
695 } else {
696 writeln!(output, "{}", insertion)?;
697 }
698 last_line_ix = insertion_line_ix + 1;
699 excerpt_file_insertions.remove(insertion_ix);
700 continue;
701 }
702 insertion_ix += 1;
703 }
704 skipped_last_snippet = false;
705 for line_ix in last_line_ix..lines.len() {
706 push_line(output, line_ix)?;
707 }
708 } else {
709 skipped_last_snippet = false;
710 push_full_snippet(output)?;
711 }
712 } else {
713 skipped_last_snippet = false;
714 push_full_snippet(output)?;
715 }
716
717 section_ranges.push((snippet.path.clone(), range));
718 }
719
720 output.push_str("`````\n\n");
721 }
722
723 Ok(SectionLabels {
724 // TODO: Clean this up
725 excerpt_index: match self.request.prompt_format {
726 PromptFormat::OnlySnippets => 0,
727 _ => excerpt_index.context("bug: no snippet found for excerpt")?,
728 },
729 section_ranges,
730 })
731 }
732}
733
734fn declaration_score_density(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> f32 {
735 declaration_score(declaration, style) / declaration_size(declaration, style) as f32
736}
737
738fn declaration_score(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> f32 {
739 match style {
740 DeclarationStyle::Signature => declaration.signature_score,
741 DeclarationStyle::Declaration => declaration.declaration_score,
742 }
743}
744
745fn declaration_size(declaration: &ReferencedDeclaration, style: DeclarationStyle) -> usize {
746 match style {
747 DeclarationStyle::Signature => declaration.signature_range.len(),
748 DeclarationStyle::Declaration => declaration.text.len(),
749 }
750}