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