From 587fd707ba9c19fcef18b6bb0f5507fab79641d9 Mon Sep 17 00:00:00 2001
From: KCaverly <kyle@zed.dev>
Date: Wed, 18 Oct 2023 16:40:09 -0400
Subject: [PATCH] added smarter error handling for file_context prompts without
 provided buffers

---
 crates/ai/src/templates/file_context.rs | 49 +++++++++++++------------
 1 file changed, 26 insertions(+), 23 deletions(-)

diff --git a/crates/ai/src/templates/file_context.rs b/crates/ai/src/templates/file_context.rs
index 00fe99dd7ffc257339caa8c5198e532b967bee40..5a6489a00c89f637f03d5c25a357f82f24accfd3 100644
--- a/crates/ai/src/templates/file_context.rs
+++ b/crates/ai/src/templates/file_context.rs
@@ -1,3 +1,4 @@
+use anyhow::anyhow;
 use language::ToOffset;
 
 use crate::templates::base::PromptArguments;
@@ -12,24 +13,23 @@ impl PromptTemplate for FileContext {
         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();
+        if let Some(buffer) = &args.buffer {
+            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();
+            let language_name = args
+                .language_name
+                .clone()
+                .unwrap_or("".to_string())
+                .to_lowercase();
+            writeln!(prompt, "```{language_name}").unwrap();
 
-        if let Some(buffer) = &args.buffer {
             if let Some(selected_range) = &args.selected_range {
                 let start = selected_range.start.to_offset(buffer);
                 let end = selected_range.end.to_offset(buffer);
@@ -74,15 +74,18 @@ impl PromptTemplate for FileContext {
             } else {
                 // If we dont have a selected range, include entire file.
                 writeln!(prompt, "{}", &buffer.text()).unwrap();
+                writeln!(prompt, "```").unwrap();
             }
-        }
 
-        // Really dumb truncation strategy
-        if let Some(max_tokens) = max_token_length {
-            prompt = args.model.truncate(&prompt, max_tokens)?;
-        }
+            // Really dumb truncation strategy
+            if let Some(max_tokens) = max_token_length {
+                prompt = args.model.truncate(&prompt, max_tokens)?;
+            }
 
-        let token_count = args.model.count_tokens(&prompt)?;
-        anyhow::Ok((prompt, token_count))
+            let token_count = args.model.count_tokens(&prompt)?;
+            anyhow::Ok((prompt, token_count))
+        } else {
+            Err(anyhow!("no buffer provided to retrieve file context from"))
+        }
     }
 }