added prompt template for file context without truncation

KCaverly created

Change summary

crates/ai/src/templates/base.rs         | 13 ++++
crates/ai/src/templates/file_context.rs | 85 ++++++++++++++++++++++++++
crates/ai/src/templates/mod.rs          |  1 
3 files changed, 99 insertions(+)

Detailed changes

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

@@ -1,6 +1,9 @@
 use std::cmp::Reverse;
+use std::ops::Range;
 use std::sync::Arc;
 
+use gpui::ModelHandle;
+use language::{Anchor, Buffer, BufferSnapshot, ToOffset};
 use util::ResultExt;
 
 use crate::models::LanguageModel;
@@ -18,6 +21,8 @@ pub struct PromptArguments {
     pub project_name: Option<String>,
     pub snippets: Vec<PromptCodeSnippet>,
     pub reserved_tokens: usize,
+    pub buffer: Option<BufferSnapshot>,
+    pub selected_range: Option<Range<usize>>,
 }
 
 impl PromptArguments {
@@ -189,6 +194,8 @@ pub(crate) mod tests {
             project_name: None,
             snippets: Vec::new(),
             reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
         };
 
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
@@ -216,6 +223,8 @@ pub(crate) mod tests {
             project_name: None,
             snippets: Vec::new(),
             reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
         };
 
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
@@ -244,6 +253,8 @@ pub(crate) mod tests {
             project_name: None,
             snippets: Vec::new(),
             reserved_tokens: 0,
+            buffer: None,
+            selected_range: None,
         };
 
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
@@ -268,6 +279,8 @@ pub(crate) mod tests {
             project_name: None,
             snippets: Vec::new(),
             reserved_tokens,
+            buffer: None,
+            selected_range: None,
         };
         let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
             (PromptPriority::Medium, Box::new(TestPromptTemplate {})),

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

@@ -0,0 +1,85 @@
+use language::ToOffset;
+
+use crate::templates::base::PromptArguments;
+use crate::templates::base::PromptTemplate;
+use std::fmt::Write;
+
+pub struct FileContext {}
+
+impl PromptTemplate for FileContext {
+    fn generate(
+        &self,
+        args: &PromptArguments,
+        max_token_length: Option<usize>,
+    ) -> anyhow::Result<(String, usize)> {
+        let mut prompt = String::new();
+
+        // Add Initial Preamble
+        // TODO: Do we want to add the path in here?
+        writeln!(
+            prompt,
+            "The file you are currently working on has the following content:"
+        )
+        .unwrap();
+
+        let language_name = args
+            .language_name
+            .clone()
+            .unwrap_or("".to_string())
+            .to_lowercase();
+        writeln!(prompt, "```{language_name}").unwrap();
+
+        if let Some(buffer) = &args.buffer {
+            let mut content = String::new();
+
+            if let Some(selected_range) = &args.selected_range {
+                let start = selected_range.start.to_offset(buffer);
+                let end = selected_range.end.to_offset(buffer);
+
+                writeln!(
+                    prompt,
+                    "{}",
+                    buffer.text_for_range(0..start).collect::<String>()
+                )
+                .unwrap();
+
+                if start == end {
+                    writeln!(prompt, "<|START|>").unwrap();
+                } else {
+                    writeln!(prompt, "<|START|").unwrap();
+                }
+
+                writeln!(
+                    prompt,
+                    "{}",
+                    buffer.text_for_range(start..end).collect::<String>()
+                )
+                .unwrap();
+                if start != end {
+                    writeln!(prompt, "|END|>").unwrap();
+                }
+
+                writeln!(
+                    prompt,
+                    "{}",
+                    buffer.text_for_range(end..buffer.len()).collect::<String>()
+                )
+                .unwrap();
+
+                writeln!(prompt, "```").unwrap();
+
+                if start == end {
+                    writeln!(prompt, "In particular, the user's cursor is currently on the '<|START|>' span in the above content, with no text selected.").unwrap();
+                } else {
+                    writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap();
+                }
+            } else {
+                // If we dont have a selected range, include entire file.
+                writeln!(prompt, "{}", &buffer.text()).unwrap();
+            }
+        }
+
+        let token_count = args.model.count_tokens(&prompt)?;
+        anyhow::Ok((prompt, token_count))
+    }
+}