prompts.rs

  1use anyhow::Result;
  2use assets::Assets;
  3use fs::Fs;
  4use futures::StreamExt;
  5use gpui::{App, AppContext as _, AssetSource};
  6use handlebars::{Handlebars, RenderError};
  7use language::{BufferSnapshot, LanguageName, Point};
  8use parking_lot::Mutex;
  9use serde::Serialize;
 10use std::{ops::Range, path::PathBuf, sync::Arc, time::Duration};
 11use text::LineEnding;
 12use util::ResultExt;
 13
 14#[derive(Serialize)]
 15pub struct AssistantSystemPromptContext {
 16    pub worktree_root_names: Vec<String>,
 17}
 18
 19#[derive(Serialize)]
 20pub struct ContentPromptDiagnosticContext {
 21    pub line_number: usize,
 22    pub error_message: String,
 23    pub code_content: String,
 24}
 25
 26#[derive(Serialize)]
 27pub struct ContentPromptContext {
 28    pub content_type: String,
 29    pub language_name: Option<String>,
 30    pub is_insert: bool,
 31    pub is_truncated: bool,
 32    pub document_content: String,
 33    pub user_prompt: String,
 34    pub rewrite_section: Option<String>,
 35    pub diagnostic_errors: Vec<ContentPromptDiagnosticContext>,
 36}
 37
 38#[derive(Serialize)]
 39pub struct TerminalAssistantPromptContext {
 40    pub os: String,
 41    pub arch: String,
 42    pub shell: Option<String>,
 43    pub working_directory: Option<String>,
 44    pub latest_output: Vec<String>,
 45    pub user_prompt: String,
 46}
 47
 48#[derive(Serialize)]
 49pub struct ProjectSlashCommandPromptContext {
 50    pub context_buffer: String,
 51}
 52
 53pub struct PromptLoadingParams<'a> {
 54    pub fs: Arc<dyn Fs>,
 55    pub repo_path: Option<PathBuf>,
 56    pub cx: &'a gpui::App,
 57}
 58
 59pub struct PromptBuilder {
 60    handlebars: Arc<Mutex<Handlebars<'static>>>,
 61}
 62
 63impl PromptBuilder {
 64    pub fn load(fs: Arc<dyn Fs>, stdout_is_a_pty: bool, cx: &mut App) -> Arc<Self> {
 65        Self::new(Some(PromptLoadingParams {
 66            fs: fs.clone(),
 67            repo_path: stdout_is_a_pty
 68                .then(|| std::env::current_dir().log_err())
 69                .flatten(),
 70            cx,
 71        }))
 72        .log_err()
 73        .map(Arc::new)
 74        .unwrap_or_else(|| Arc::new(Self::new(None).unwrap()))
 75    }
 76
 77    pub fn new(loading_params: Option<PromptLoadingParams>) -> Result<Self> {
 78        let mut handlebars = Handlebars::new();
 79        Self::register_built_in_templates(&mut handlebars)?;
 80
 81        let handlebars = Arc::new(Mutex::new(handlebars));
 82
 83        if let Some(params) = loading_params {
 84            Self::watch_fs_for_template_overrides(params, handlebars.clone());
 85        }
 86
 87        Ok(Self { handlebars })
 88    }
 89
 90    /// Watches the filesystem for changes to prompt template overrides.
 91    ///
 92    /// This function sets up a file watcher on the prompt templates directory. It performs
 93    /// an initial scan of the directory and registers any existing template overrides.
 94    /// Then it continuously monitors for changes, reloading templates as they are
 95    /// modified or added.
 96    ///
 97    /// If the templates directory doesn't exist initially, it waits for it to be created.
 98    /// If the directory is removed, it restores the built-in templates and waits for the
 99    /// directory to be recreated.
100    ///
101    /// # Arguments
102    ///
103    /// * `params` - A `PromptLoadingParams` struct containing the filesystem, repository path,
104    ///   and application context.
105    /// * `handlebars` - An `Arc<Mutex<Handlebars>>` for registering and updating templates.
106    fn watch_fs_for_template_overrides(
107        params: PromptLoadingParams,
108        handlebars: Arc<Mutex<Handlebars<'static>>>,
109    ) {
110        let templates_dir = paths::prompt_overrides_dir(params.repo_path.as_deref());
111        params.cx.background_spawn(async move {
112            let Some(parent_dir) = templates_dir.parent() else {
113                return;
114            };
115
116            let mut found_dir_once = false;
117            loop {
118                // Check if the templates directory exists and handle its status
119                // If it exists, log its presence and check if it's a symlink
120                // If it doesn't exist:
121                //   - Log that we're using built-in prompts
122                //   - Check if it's a broken symlink and log if so
123                //   - Set up a watcher to detect when it's created
124                // After the first check, set the `found_dir_once` flag
125                // This allows us to avoid logging when looping back around after deleting the prompt overrides directory.
126                let dir_status = params.fs.is_dir(&templates_dir).await;
127                let symlink_status = params.fs.read_link(&templates_dir).await.ok();
128                if dir_status {
129                    let mut log_message = format!("Prompt template overrides directory found at {}", templates_dir.display());
130                    if let Some(target) = symlink_status {
131                        log_message.push_str(" -> ");
132                        log_message.push_str(&target.display().to_string());
133                    }
134                    log::info!("{}.", log_message);
135                } else {
136                    if !found_dir_once {
137                        log::info!("No prompt template overrides directory found at {}. Using built-in prompts.", templates_dir.display());
138                        if let Some(target) = symlink_status {
139                            log::info!("Symlink found pointing to {}, but target is invalid.", target.display());
140                        }
141                    }
142
143                    if params.fs.is_dir(parent_dir).await {
144                        let (mut changes, _watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
145                        while let Some(changed_paths) = changes.next().await {
146                            if changed_paths.iter().any(|p| &p.path == &templates_dir) {
147                                let mut log_message = format!("Prompt template overrides directory detected at {}", templates_dir.display());
148                                if let Ok(target) = params.fs.read_link(&templates_dir).await {
149                                    log_message.push_str(" -> ");
150                                    log_message.push_str(&target.display().to_string());
151                                }
152                                log::info!("{}.", log_message);
153                                break;
154                            }
155                        }
156                    } else {
157                        return;
158                    }
159                }
160
161                found_dir_once = true;
162
163                // Initial scan of the prompt overrides directory
164                if let Ok(mut entries) = params.fs.read_dir(&templates_dir).await {
165                    while let Some(Ok(file_path)) = entries.next().await {
166                        if file_path.to_string_lossy().ends_with(".hbs") {
167                            if let Ok(content) = params.fs.load(&file_path).await {
168                                let file_name = file_path.file_stem().unwrap().to_string_lossy();
169                                log::debug!("Registering prompt template override: {}", file_name);
170                                handlebars.lock().register_template_string(&file_name, content).log_err();
171                            }
172                        }
173                    }
174                }
175
176                // Watch both the parent directory and the template overrides directory:
177                // - Monitor the parent directory to detect if the template overrides directory is deleted.
178                // - Monitor the template overrides directory to re-register templates when they change.
179                // Combine both watch streams into a single stream.
180                let (parent_changes, parent_watcher) = params.fs.watch(parent_dir, Duration::from_secs(1)).await;
181                let (changes, watcher) = params.fs.watch(&templates_dir, Duration::from_secs(1)).await;
182                let mut combined_changes = futures::stream::select(changes, parent_changes);
183
184                while let Some(changed_paths) = combined_changes.next().await {
185                    if changed_paths.iter().any(|p| &p.path == &templates_dir) {
186                        if !params.fs.is_dir(&templates_dir).await {
187                            log::info!("Prompt template overrides directory removed. Restoring built-in prompt templates.");
188                            Self::register_built_in_templates(&mut handlebars.lock()).log_err();
189                            break;
190                        }
191                    }
192                    for event in changed_paths {
193                        if event.path.starts_with(&templates_dir) && event.path.extension().map_or(false, |ext| ext == "hbs") {
194                            log::info!("Reloading prompt template override: {}", event.path.display());
195                            if let Some(content) = params.fs.load(&event.path).await.log_err() {
196                                let file_name = event.path.file_stem().unwrap().to_string_lossy();
197                                handlebars.lock().register_template_string(&file_name, content).log_err();
198                            }
199                        }
200                    }
201                }
202
203                drop(watcher);
204                drop(parent_watcher);
205            }
206        })
207            .detach();
208    }
209
210    fn register_built_in_templates(handlebars: &mut Handlebars) -> Result<()> {
211        for path in Assets.list("prompts")? {
212            if let Some(id) = path.split('/').last().and_then(|s| s.strip_suffix(".hbs")) {
213                if let Some(prompt) = Assets.load(path.as_ref()).log_err().flatten() {
214                    log::debug!("Registering built-in prompt template: {}", id);
215                    let prompt = String::from_utf8_lossy(prompt.as_ref());
216                    handlebars.register_template_string(id, LineEnding::normalize_cow(prompt))?
217                }
218            }
219        }
220
221        Ok(())
222    }
223
224    pub fn generate_assistant_system_prompt(
225        &self,
226        worktree_root_names: Vec<String>,
227    ) -> Result<String, RenderError> {
228        let prompt = AssistantSystemPromptContext {
229            worktree_root_names,
230        };
231        self.handlebars
232            .lock()
233            .render("assistant_system_prompt", &prompt)
234    }
235
236    pub fn generate_inline_transformation_prompt(
237        &self,
238        user_prompt: String,
239        language_name: Option<&LanguageName>,
240        buffer: BufferSnapshot,
241        range: Range<usize>,
242    ) -> Result<String, RenderError> {
243        let content_type = match language_name.as_ref().map(|l| l.as_ref()) {
244            None | Some("Markdown" | "Plain Text") => "text",
245            Some(_) => "code",
246        };
247
248        const MAX_CTX: usize = 50000;
249        let is_insert = range.is_empty();
250        let mut is_truncated = false;
251
252        let before_range = 0..range.start;
253        let truncated_before = if before_range.len() > MAX_CTX {
254            is_truncated = true;
255            let start = buffer.clip_offset(range.start - MAX_CTX, text::Bias::Right);
256            start..range.start
257        } else {
258            before_range
259        };
260
261        let after_range = range.end..buffer.len();
262        let truncated_after = if after_range.len() > MAX_CTX {
263            is_truncated = true;
264            let end = buffer.clip_offset(range.end + MAX_CTX, text::Bias::Left);
265            range.end..end
266        } else {
267            after_range
268        };
269
270        let mut document_content = String::new();
271        for chunk in buffer.text_for_range(truncated_before) {
272            document_content.push_str(chunk);
273        }
274        if is_insert {
275            document_content.push_str("<insert_here></insert_here>");
276        } else {
277            document_content.push_str("<rewrite_this>\n");
278            for chunk in buffer.text_for_range(range.clone()) {
279                document_content.push_str(chunk);
280            }
281            document_content.push_str("\n</rewrite_this>");
282        }
283        for chunk in buffer.text_for_range(truncated_after) {
284            document_content.push_str(chunk);
285        }
286
287        let rewrite_section = if !is_insert {
288            let mut section = String::new();
289            for chunk in buffer.text_for_range(range.clone()) {
290                section.push_str(chunk);
291            }
292            Some(section)
293        } else {
294            None
295        };
296        let diagnostics = buffer.diagnostics_in_range::<_, Point>(range, false);
297        let diagnostic_errors: Vec<ContentPromptDiagnosticContext> = diagnostics
298            .map(|entry| {
299                let start = entry.range.start;
300                ContentPromptDiagnosticContext {
301                    line_number: (start.row + 1) as usize,
302                    error_message: entry.diagnostic.message.clone(),
303                    code_content: buffer.text_for_range(entry.range.clone()).collect(),
304                }
305            })
306            .collect();
307
308        let context = ContentPromptContext {
309            content_type: content_type.to_string(),
310            language_name: language_name.map(|s| s.to_string()),
311            is_insert,
312            is_truncated,
313            document_content,
314            user_prompt,
315            rewrite_section,
316            diagnostic_errors,
317        };
318        self.handlebars.lock().render("content_prompt", &context)
319    }
320
321    pub fn generate_terminal_assistant_prompt(
322        &self,
323        user_prompt: &str,
324        shell: Option<&str>,
325        working_directory: Option<&str>,
326        latest_output: &[String],
327    ) -> Result<String, RenderError> {
328        let context = TerminalAssistantPromptContext {
329            os: std::env::consts::OS.to_string(),
330            arch: std::env::consts::ARCH.to_string(),
331            shell: shell.map(|s| s.to_string()),
332            working_directory: working_directory.map(|s| s.to_string()),
333            latest_output: latest_output.to_vec(),
334            user_prompt: user_prompt.to_string(),
335        };
336
337        self.handlebars
338            .lock()
339            .render("terminal_assistant_prompt", &context)
340    }
341
342    pub fn generate_suggest_edits_prompt(&self) -> Result<String, RenderError> {
343        self.handlebars.lock().render("suggest_edits", &())
344    }
345
346    pub fn generate_project_slash_command_prompt(
347        &self,
348        context_buffer: String,
349    ) -> Result<String, RenderError> {
350        self.handlebars.lock().render(
351            "project_slash_command",
352            &ProjectSlashCommandPromptContext { context_buffer },
353        )
354    }
355}