1use crate::templates::base::{PromptArguments, PromptFileType, PromptTemplate};
2use anyhow::anyhow;
3use std::fmt::Write;
4
5pub fn capitalize(s: &str) -> String {
6 let mut c = s.chars();
7 match c.next() {
8 None => String::new(),
9 Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
10 }
11}
12
13pub struct GenerateInlineContent {}
14
15impl PromptTemplate for GenerateInlineContent {
16 fn generate(
17 &self,
18 args: &PromptArguments,
19 max_token_length: Option<usize>,
20 ) -> anyhow::Result<(String, usize)> {
21 let Some(user_prompt) = &args.user_prompt else {
22 return Err(anyhow!("user prompt not provided"));
23 };
24
25 let file_type = args.get_file_type();
26 let content_type = match &file_type {
27 PromptFileType::Code => "code",
28 PromptFileType::Text => "text",
29 };
30
31 let mut prompt = String::new();
32
33 if let Some(selected_range) = &args.selected_range {
34 if selected_range.start == selected_range.end {
35 writeln!(
36 prompt,
37 "Assume the cursor is located where the `<|START|>` span is."
38 )
39 .unwrap();
40 writeln!(
41 prompt,
42 "{} can't be replaced, so assume your answer will be inserted at the cursor.",
43 capitalize(content_type)
44 )
45 .unwrap();
46 writeln!(
47 prompt,
48 "Generate {content_type} based on the users prompt: {user_prompt}",
49 )
50 .unwrap();
51 } else {
52 writeln!(prompt, "Modify the user's selected {content_type} based upon the users prompt: '{user_prompt}'").unwrap();
53 writeln!(prompt, "You MUST reply with only the adjusted {content_type} (within the '<|START|' and '|END|>' spans), not the entire file.").unwrap();
54 }
55 } else {
56 writeln!(
57 prompt,
58 "Generate {content_type} based on the users prompt: {user_prompt}"
59 )
60 .unwrap();
61 }
62
63 if let Some(language_name) = &args.language_name {
64 writeln!(
65 prompt,
66 "Your answer MUST always and only be valid {}.",
67 language_name
68 )
69 .unwrap();
70 }
71 writeln!(prompt, "Never make remarks about the output.").unwrap();
72 writeln!(
73 prompt,
74 "Do not return anything else, except the generated {content_type}."
75 )
76 .unwrap();
77
78 match file_type {
79 PromptFileType::Code => {
80 writeln!(prompt, "Always wrap your code in a Markdown block.").unwrap();
81 }
82 _ => {}
83 }
84
85 // Really dumb truncation strategy
86 if let Some(max_tokens) = max_token_length {
87 prompt = args.model.truncate(&prompt, max_tokens)?;
88 }
89
90 let token_count = args.model.count_tokens(&prompt)?;
91
92 anyhow::Ok((prompt, token_count))
93 }
94}