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