zed-industries/zed
1
1
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2024-12-28 22:24:47 +03:00
Code Issues Packages Projects Releases Wiki Activity

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
Show Stats 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 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();
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 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 {
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"))
}
}
}