prompts.rs

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