1use crate::prompts::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 writeln!(prompt, "Double check that you only return code and not the '<|START|' and '|END|'> spans").unwrap();
55 }
56 } else {
57 writeln!(
58 prompt,
59 "Generate {content_type} based on the users prompt: {user_prompt}"
60 )
61 .unwrap();
62 }
63
64 if let Some(language_name) = &args.language_name {
65 writeln!(
66 prompt,
67 "Your answer MUST always and only be valid {}.",
68 language_name
69 )
70 .unwrap();
71 }
72 writeln!(prompt, "Never make remarks about the output.").unwrap();
73 writeln!(
74 prompt,
75 "Do not return anything else, except the generated {content_type}."
76 )
77 .unwrap();
78
79 match file_type {
80 PromptFileType::Code => {
81 // writeln!(prompt, "Always wrap your code in a Markdown block.").unwrap();
82 }
83 _ => {}
84 }
85
86 // Really dumb truncation strategy
87 if let Some(max_tokens) = max_token_length {
88 prompt = args.model.truncate(
89 &prompt,
90 max_tokens,
91 crate::models::TruncationDirection::End,
92 )?;
93 }
94
95 let token_count = args.model.count_tokens(&prompt)?;
96
97 anyhow::Ok((prompt, token_count))
98 }
99}