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

Allow prompt templates to be overridden in the zed configuration directory (#15887)

Browse source 
I need this to refine our prompts on the fly as I work.

Release Notes:

- Templates for prompts driving inline transformation in editors and the
terminal can now be overridden in the `~/.config/zed/prompts/templates`
directory. This is an advanced feature, and prevents you from getting
upstream changes. It's intended for use by Zed developers.
This commit is contained in:
Nathan Sobo 2024-08-06 19:30:48 -06:00 • committed by GitHub
parent 6065db174a
commit c8f1358629
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 569 additions and 168 deletions
Download patch file Download diff file Expand all files Collapse all files

364
crates/assistant/src/prompts.rs
View file

@ -1,153 +1,239 @@
use assets::Assets;
use fs::Fs;
use futures::StreamExt;
use handlebars::{Handlebars, RenderError, TemplateError};
use language::BufferSnapshot;
use std::{fmt::Write, ops::Range};
use parking_lot::Mutex;
use serde::Serialize;
use std::{ops::Range, sync::Arc, time::Duration};
use util::ResultExt;
pub fn generate_content_prompt(
user_prompt: String,
language_name: Option<&str>,
buffer: BufferSnapshot,
range: Range<usize>,
) -> String {
let mut prompt = String::new();
#[derive(Serialize)]
pub struct ContentPromptContext {
pub content_type: String,
pub language_name: Option<String>,
pub is_insert: bool,
pub is_truncated: bool,
pub document_content: String,
pub user_prompt: String,
pub rewrite_section: Option<String>,
}
let content_type = match language_name {
None | Some("Markdown" | "Plain Text") => {
writeln!(
prompt,
"Here's a file of text that I'm going to ask you to make an edit to."
)
.unwrap();
"text"
#[derive(Serialize)]
pub struct TerminalAssistantPromptContext {
pub os: String,
pub arch: String,
pub shell: Option<String>,
pub working_directory: Option<String>,
pub latest_output: Vec<String>,
pub user_prompt: String,
}
pub struct PromptBuilder {
handlebars: Arc<Mutex<Handlebars<'static>>>,
}
impl PromptBuilder {
pub fn new(
fs_and_cx: Option<(Arc<dyn Fs>, &gpui::AppContext)>,
) -> Result<Self, Box<TemplateError>> {
let mut handlebars = Handlebars::new();
Self::register_templates(&mut handlebars)?;
let handlebars = Arc::new(Mutex::new(handlebars));
if let Some((fs, cx)) = fs_and_cx {
Self::watch_fs_for_template_overrides(fs, cx, handlebars.clone());
}
Some(language_name) => {
writeln!(
prompt,
"Here's a file of {language_name} that I'm going to ask you to make an edit to."
Ok(Self { handlebars })
}
fn watch_fs_for_template_overrides(
fs: Arc<dyn Fs>,
cx: &gpui::AppContext,
handlebars: Arc<Mutex<Handlebars<'static>>>,
) {
let templates_dir = paths::prompt_templates_dir();
cx.background_executor()
.spawn(async move {
// Create the prompt templates directory if it doesn't exist
if !fs.is_dir(templates_dir).await {
if let Err(e) = fs.create_dir(templates_dir).await {
log::error!("Failed to create prompt templates directory: {}", e);
return;
}
}
// Initial scan of the prompts directory
if let Ok(mut entries) = fs.read_dir(templates_dir).await {
while let Some(Ok(file_path)) = entries.next().await {
if file_path.to_string_lossy().ends_with(".hbs") {
if let Ok(content) = fs.load(&file_path).await {
let file_name = file_path.file_stem().unwrap().to_string_lossy();
match handlebars.lock().register_template_string(&file_name, content) {
Ok(_) => {
log::info!(
"Successfully registered template override: {} ({})",
file_name,
file_path.display()
);
},
Err(e) => {
log::error!(
"Failed to register template during initial scan: {} ({})",
e,
file_path.display()
);
},
}
}
}
}
}
// Watch for changes
let (mut changes, watcher) = fs.watch(templates_dir, Duration::from_secs(1)).await;
while let Some(changed_paths) = changes.next().await {
for changed_path in changed_paths {
if changed_path.extension().map_or(false, |ext| ext == "hbs") {
log::info!("Reloading template: {}", changed_path.display());
if let Some(content) = fs.load(&changed_path).await.log_err() {
let file_name = changed_path.file_stem().unwrap().to_string_lossy();
let file_path = changed_path.to_string_lossy();
match handlebars.lock().register_template_string(&file_name, content) {
Ok(_) => log::info!(
"Successfully reloaded template: {} ({})",
file_name,
file_path
),
Err(e) => log::error!(
"Failed to register template: {} ({})",
e,
file_path
),
}
}
}
}
}
drop(watcher);
})
.detach();
}
fn register_templates(handlebars: &mut Handlebars) -> Result<(), Box<TemplateError>> {
let content_prompt = Assets::get("prompts/content_prompt.hbs")
.expect("Content prompt template not found")
.data;
let terminal_assistant_prompt = Assets::get("prompts/terminal_assistant_prompt.hbs")
.expect("Terminal assistant prompt template not found")
.data;
handlebars
.register_template_string("content_prompt", String::from_utf8_lossy(&content_prompt))
.map_err(Box::new)?;
handlebars
.register_template_string(
"terminal_assistant_prompt",
String::from_utf8_lossy(&terminal_assistant_prompt),
)
.unwrap();
"code"
.map_err(Box::new)?;
Ok(())
}
pub fn generate_content_prompt(
&self,
user_prompt: String,
language_name: Option<&str>,
buffer: BufferSnapshot,
range: Range<usize>,
) -> Result<String, RenderError> {
let content_type = match language_name {
None | Some("Markdown" | "Plain Text") => "text",
Some(_) => "code",
};
const MAX_CTX: usize = 50000;
let is_insert = range.is_empty();
let mut is_truncated = false;
let before_range = 0..range.start;
let truncated_before = if before_range.len() > MAX_CTX {
is_truncated = true;
range.start - MAX_CTX..range.start
} else {
before_range
};
let after_range = range.end..buffer.len();
let truncated_after = if after_range.len() > MAX_CTX {
is_truncated = true;
range.end..range.end + MAX_CTX
} else {
after_range
};
let mut document_content = String::new();
for chunk in buffer.text_for_range(truncated_before) {
document_content.push_str(chunk);
}
};
const MAX_CTX: usize = 50000;
let mut is_truncated = false;
if range.is_empty() {
prompt.push_str("The point you'll need to insert at is marked with <insert_here></insert_here>.\n\n<document>");
} else {
prompt.push_str("The section you'll need to rewrite is marked with <rewrite_this></rewrite_this> tags.\n\n<document>");
}
// Include file content.
let before_range = 0..range.start;
let truncated_before = if before_range.len() > MAX_CTX {
is_truncated = true;
range.start - MAX_CTX..range.start
} else {
before_range
};
let mut non_rewrite_len = truncated_before.len();
for chunk in buffer.text_for_range(truncated_before) {
prompt.push_str(chunk);
}
if !range.is_empty() {
prompt.push_str("<rewrite_this>\n");
for chunk in buffer.text_for_range(range.clone()) {
prompt.push_str(chunk);
}
prompt.push_str("\n<rewrite_this>");
} else {
prompt.push_str("<insert_here></insert_here>");
}
let after_range = range.end..buffer.len();
let truncated_after = if after_range.len() > MAX_CTX {
is_truncated = true;
range.end..range.end + MAX_CTX
} else {
after_range
};
non_rewrite_len += truncated_after.len();
for chunk in buffer.text_for_range(truncated_after) {
prompt.push_str(chunk);
}
write!(prompt, "</document>\n\n").unwrap();
if is_truncated {
writeln!(prompt, "The context around the relevant section has been truncated (possibly in the middle of a line) for brevity.\n").unwrap();
}
if range.is_empty() {
writeln!(
prompt,
"You can't replace {content_type}, your answer will be inserted in place of the `<insert_here></insert_here>` tags. Don't include the insert_here tags in your output.",
)
.unwrap();
writeln!(
prompt,
"Generate {content_type} based on the following prompt:\n\n<prompt>\n{user_prompt}\n</prompt>",
)
.unwrap();
writeln!(prompt, "Match the indentation in the original file in the inserted {content_type}, don't include any indentation on blank lines.\n").unwrap();
prompt.push_str("Immediately start with the following format with no remarks:\n\n```\n{{INSERTED_CODE}}\n```");
} else {
writeln!(prompt, "Edit the section of {content_type} in <rewrite_this></rewrite_this> tags based on the following prompt:'").unwrap();
writeln!(prompt, "\n<prompt>\n{user_prompt}\n</prompt>\n").unwrap();
let rewrite_len = range.end - range.start;
if rewrite_len < 20000 && rewrite_len * 2 < non_rewrite_len {
writeln!(prompt, "And here's the section to rewrite based on that prompt again for reference:\n\n<rewrite_this>\n").unwrap();
if is_insert {
document_content.push_str("<insert_here></insert_here>");
} else {
document_content.push_str("<rewrite_this>\n");
for chunk in buffer.text_for_range(range.clone()) {
prompt.push_str(chunk);
document_content.push_str(chunk);
}
writeln!(prompt, "\n</rewrite_this>\n").unwrap();
document_content.push_str("\n</rewrite_this>");
}
writeln!(prompt, "Only make changes that are necessary to fulfill the prompt, leave everything else as-is. All surrounding {content_type} will be preserved.\n").unwrap();
write!(
prompt,
"Start at the indentation level in the original file in the rewritten {content_type}. "
)
.unwrap();
prompt.push_str("Don't stop until you've rewritten the entire section, even if you have no more changes to make, always write out the whole section with no unnecessary elisions.");
prompt.push_str("\n\nImmediately start with the following format with no remarks:\n\n```\n{{REWRITTEN_CODE}}\n```");
for chunk in buffer.text_for_range(truncated_after) {
document_content.push_str(chunk);
}
let rewrite_section = if !is_insert {
let mut section = String::new();
for chunk in buffer.text_for_range(range.clone()) {
section.push_str(chunk);
}
Some(section)
} else {
None
};
let context = ContentPromptContext {
content_type: content_type.to_string(),
language_name: language_name.map(|s| s.to_string()),
is_insert,
is_truncated,
document_content,
user_prompt,
rewrite_section,
};
self.handlebars.lock().render("content_prompt", &context)
}
prompt
}
pub fn generate_terminal_assistant_prompt(
&self,
user_prompt: &str,
shell: Option<&str>,
working_directory: Option<&str>,
latest_output: &[String],
) -> Result<String, RenderError> {
let context = TerminalAssistantPromptContext {
os: std::env::consts::OS.to_string(),
arch: std::env::consts::ARCH.to_string(),
shell: shell.map(|s| s.to_string()),
working_directory: working_directory.map(|s| s.to_string()),
latest_output: latest_output.to_vec(),
user_prompt: user_prompt.to_string(),
};
pub fn generate_terminal_assistant_prompt(
user_prompt: &str,
shell: Option<&str>,
working_directory: Option<&str>,
latest_output: &[String],
) -> String {
let mut prompt = String::new();
writeln!(&mut prompt, "You are an expert terminal user.").unwrap();
writeln!(&mut prompt, "You will be given a description of a command and you need to respond with a command that matches the description.").unwrap();
writeln!(&mut prompt, "Do not include markdown blocks or any other text formatting in your response, always respond with a single command that can be executed in the given shell.").unwrap();
writeln!(
&mut prompt,
"Current OS name is '{}', architecture is '{}'.",
std::env::consts::OS,
std::env::consts::ARCH,
)
.unwrap();
if let Some(shell) = shell {
writeln!(&mut prompt, "Current shell is '{shell}'.").unwrap();
self.handlebars
.lock()
.render("terminal_assistant_prompt", &context)
}
if let Some(working_directory) = working_directory {
writeln!(
&mut prompt,
"Current working directory is '{working_directory}'."
)
.unwrap();
}
if !latest_output.is_empty() {
writeln!(
&mut prompt,
"Latest non-empty {} lines of the terminal output: {:?}",
latest_output.len(),
latest_output
)
.unwrap();
}
writeln!(&mut prompt, "Here is the description of the command:").unwrap();
prompt.push_str(user_prompt);
prompt
}