Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

added smarter error handling for file_context prompts without provided buffers

Browse source
This commit is contained in:
KCaverly 2023-10-18 16:40:09 -04:00
parent f59f2eccd5
commit 587fd707ba
1 changed files with 27 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

51
crates/ai/src/templates/file_context.rs
View file

@ -1,3 +1,4 @@
use anyhow::anyhow;
use language::ToOffset; use language::ToOffset;
use crate::templates::base::PromptArguments; use crate::templates::base::PromptArguments;
@ -12,24 +13,23 @@ impl PromptTemplate for FileContext {
args: &PromptArguments, args: &PromptArguments,
max_token_length: Option<usize>, max_token_length: Option<usize>,
) -> anyhow::Result<(String, 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 { 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();
if let Some(selected_range) = &args.selected_range { if let Some(selected_range) = &args.selected_range {
let start = selected_range.start.to_offset(buffer); let start = selected_range.start.to_offset(buffer);
let end = selected_range.end.to_offset(buffer); let end = selected_range.end.to_offset(buffer);
@ -74,15 +74,18 @@ impl PromptTemplate for FileContext {
} else { } else {
// If we dont have a selected range, include entire file. // If we dont have a selected range, include entire file.
writeln!(prompt, "{}", &buffer.text()).unwrap(); writeln!(prompt, "{}", &buffer.text()).unwrap();
writeln!(prompt, "```").unwrap();
} }
}
// Really dumb truncation strategy // Really dumb truncation strategy
if let Some(max_tokens) = max_token_length { if let Some(max_tokens) = max_token_length {
prompt = args.model.truncate(&prompt, max_tokens)?; prompt = args.model.truncate(&prompt, max_tokens)?;
} }
let token_count = args.model.count_tokens(&prompt)?; let token_count = args.model.count_tokens(&prompt)?;
anyhow::Ok((prompt, token_count)) anyhow::Ok((prompt, token_count))
} else {
Err(anyhow!("no buffer provided to retrieve file context from"))
}
} }
} }