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