Unified · Split
| @@ -1,3 +1,4 @@ | |||
| 1 | use anyhow::anyhow; | ||
| 1 | use language::ToOffset; | 2 | use language::ToOffset; |
| 2 | 3 | ||
| 3 | use crate::templates::base::PromptArguments; | 4 | use crate::templates::base::PromptArguments; |
| @@ -12,24 +13,23 @@ impl PromptTemplate for FileContext { · 8 unmodified lines | |||
| 12 | args: &PromptArguments, | 13 | args: &PromptArguments, |
| 13 | max_token_length: Option<usize>, | 14 | max_token_length: Option<usize>, |
| 14 | ) -> anyhow::Result<(String, usize)> { | 15 | ) -> anyhow::Result<(String, usize)> { |
| 15 | let mut prompt = String::new(); | 16 | if let Some(buffer) = &args.buffer { |
| 16 | 17 | let mut prompt = String::new(); | |
| 17 | // Add Initial Preamble | 18 | // Add Initial Preamble |
| 18 | // TODO: Do we want to add the path in here? | 19 | // TODO: Do we want to add the path in here? |
| 19 | writeln!( | 20 | writeln!( |
| 20 | prompt, | 21 | prompt, |
| 21 | "The file you are currently working on has the following content:" | 22 | "The file you are currently working on has the following content:" |
| 22 | ) | 23 | ) |
| 23 | .unwrap(); | 24 | .unwrap(); |
| 24 | 25 | ||
| 25 | let language_name = args | 26 | let language_name = args |
| 26 | .language_name | 27 | .language_name |
| 27 | .clone() | 28 | .clone() |
| 28 | .unwrap_or("".to_string()) | 29 | .unwrap_or("".to_string()) |
| 29 | .to_lowercase(); | 30 | .to_lowercase(); |
| 30 | writeln!(prompt, "```{language_name}").unwrap(); | 31 | writeln!(prompt, "```{language_name}").unwrap(); |
| 31 | 32 | ||
| 32 | if let Some(buffer) = &args.buffer { | ||
| 33 | if let Some(selected_range) = &args.selected_range { | 33 | if let Some(selected_range) = &args.selected_range { |
| 34 | let start = selected_range.start.to_offset(buffer); | 34 | let start = selected_range.start.to_offset(buffer); |
| 35 | let end = selected_range.end.to_offset(buffer); | 35 | let end = selected_range.end.to_offset(buffer); |
| @@ -74,15 +74,18 @@ impl PromptTemplate for FileContext { · 38 unmodified lines | |||
| 74 | } else { | 74 | } else { |
| 75 | // If we dont have a selected range, include entire file. | 75 | // If we dont have a selected range, include entire file. |
| 76 | writeln!(prompt, "{}", &buffer.text()).unwrap(); | 76 | writeln!(prompt, "{}", &buffer.text()).unwrap(); |
| 77 | writeln!(prompt, "```").unwrap(); | ||
| 77 | } | 78 | } |
| 78 | } | ||
| 79 | 79 | ||
| 80 | // Really dumb truncation strategy | 80 | // Really dumb truncation strategy |
| 81 | if let Some(max_tokens) = max_token_length { | 81 | if let Some(max_tokens) = max_token_length { |
| 82 | prompt = args.model.truncate(&prompt, max_tokens)?; | 82 | prompt = args.model.truncate(&prompt, max_tokens)?; |
| 83 | } | 83 | } |
| 84 | 84 | ||
| 85 | let token_count = args.model.count_tokens(&prompt)?; | 85 | let token_count = args.model.count_tokens(&prompt)?; |
| 86 | anyhow::Ok((prompt, token_count)) | 86 | anyhow::Ok((prompt, token_count)) |
| 87 | } else { | ||
| 88 | Err(anyhow!("no buffer provided to retrieve file context from")) | ||
| 89 | } | ||
| 87 | } | 90 | } |
| 88 | } | 91 | } |