add prompt template for generate inline content

KCaverly created

Change summary

crates/ai/src/templates/base.rs     |  5 +
crates/ai/src/templates/generate.rs | 88 +++++++++++++++++++++++++++++++
crates/ai/src/templates/mod.rs      |  1 
3 files changed, 94 insertions(+)

Detailed changes

crates/ai/src/templates/base.rs 🔗 

@@ -17,6 +17,7 @@ pub(crate) enum PromptFileType {
 // TODO: Set this up to manage for defaults well
 pub struct PromptArguments {
     pub model: Arc<dyn LanguageModel>,
+    pub user_prompt: Option<String>,
     pub language_name: Option<String>,
     pub project_name: Option<String>,
     pub snippets: Vec<PromptCodeSnippet>,
@@ -196,6 +197,7 @@ pub(crate) mod tests {
             reserved_tokens: 0,
             buffer: None,
             selected_range: None,
+            user_prompt: None,
         };
 
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
@@ -225,6 +227,7 @@ pub(crate) mod tests {
             reserved_tokens: 0,
             buffer: None,
             selected_range: None,
+            user_prompt: None,
         };
 
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
@@ -255,6 +258,7 @@ pub(crate) mod tests {
             reserved_tokens: 0,
             buffer: None,
             selected_range: None,
+            user_prompt: None,
         };
 
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
@@ -281,6 +285,7 @@ pub(crate) mod tests {
             reserved_tokens,
             buffer: None,
             selected_range: None,
+            user_prompt: None,
         };
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
             (PromptPriority::Medium, Box::new(TestPromptTemplate {})),

crates/ai/src/templates/generate.rs 🔗 

@@ -0,0 +1,88 @@
+use crate::templates::base::{PromptArguments, PromptFileType, PromptTemplate};
+use anyhow::anyhow;
+use std::fmt::Write;
+
+pub fn capitalize(s: &str) -> String {
+    let mut c = s.chars();
+    match c.next() {
+        None => String::new(),
+        Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
+    }
+}
+
+pub struct GenerateInlineContent {}
+
+impl PromptTemplate for GenerateInlineContent {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        let Some(user_prompt) = &args.user_prompt else {
+            return Err(anyhow!("user prompt not provided"));
+        };
+
+        let file_type = args.get_file_type();
+        let content_type = match &file_type {
+            PromptFileType::Code => "code",
+            PromptFileType::Text => "text",
+        };
+
+        let mut prompt = String::new();
+
+        if let Some(selected_range) = &args.selected_range {
+            if selected_range.start == selected_range.end {
+                writeln!(
+                    prompt,
+                    "Assume the cursor is located where the `<|START|>` span is."
+                )
+                .unwrap();
+                writeln!(
+                    prompt,
+                    "{} can't be replaced, so assume your answer will be inserted at the cursor.",
+                    capitalize(content_type)
+                )
+                .unwrap();
+                writeln!(
+                    prompt,
+                    "Generate {content_type} based on the users prompt: {user_prompt}",
+                )
+                .unwrap();
+            } else {
+                writeln!(prompt, "Modify the user's selected {content_type} based upon the users prompt: '{user_prompt}'").unwrap();
+                writeln!(prompt, "You MUST reply with only the adjusted {content_type} (within the '<|START|' and '|END|>' spans), not the entire file.").unwrap();
+            }
+        } else {
+            writeln!(
+                prompt,
+                "Generate {content_type} based on the users prompt: {user_prompt}"
+            )
+            .unwrap();
+        }
+
+        if let Some(language_name) = &args.language_name {
+            writeln!(
+                prompt,
+                "Your answer MUST always and only be valid {}.",
+                language_name
+            )
+            .unwrap();
+        }
+        writeln!(prompt, "Never make remarks about the output.").unwrap();
+        writeln!(
+            prompt,
+            "Do not return anything else, except the generated {content_type}."
+        )
+        .unwrap();
+
+        match file_type {
+            PromptFileType::Code => {
+                writeln!(prompt, "Always wrap your code in a Markdown block.").unwrap();
+            }
+            _ => {}
+        }
+
+        let token_count = args.model.count_tokens(&prompt)?;
+        anyhow::Ok((prompt, token_count))
+    }
+}