prompt_command.rs

  1use anyhow::{Context as _, Result, anyhow};
  2use assistant_slash_command::{
  3    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,
  4    SlashCommandResult,
  5};
  6use gpui::{Task, WeakEntity};
  7use language::{BufferSnapshot, LspAdapterDelegate};
  8use prompt_store::{PromptMetadata, PromptStore};
  9use std::sync::{Arc, atomic::AtomicBool};
 10use ui::prelude::*;
 11use workspace::Workspace;
 12
 13pub struct PromptSlashCommand;
 14
 15impl SlashCommand for PromptSlashCommand {
 16    fn name(&self) -> String {
 17        "prompt".into()
 18    }
 19
 20    fn description(&self) -> String {
 21        "Insert prompt from library".into()
 22    }
 23
 24    fn icon(&self) -> IconName {
 25        IconName::Library
 26    }
 27
 28    fn menu_text(&self) -> String {
 29        self.description()
 30    }
 31
 32    fn requires_argument(&self) -> bool {
 33        true
 34    }
 35
 36    fn complete_argument(
 37        self: Arc<Self>,
 38        arguments: &[String],
 39        _cancellation_flag: Arc<AtomicBool>,
 40        _workspace: Option<WeakEntity<Workspace>>,
 41        _: &mut Window,
 42        cx: &mut App,
 43    ) -> Task<Result<Vec<ArgumentCompletion>>> {
 44        let store = PromptStore::global(cx);
 45        let query = arguments.to_owned().join(" ");
 46        cx.spawn(async move |cx| {
 47            let prompts: Vec<PromptMetadata> = store
 48                .await?
 49                .read_with(cx, |store, cx| store.search(query, cx))?
 50                .await;
 51            Ok(prompts
 52                .into_iter()
 53                .filter_map(|prompt| {
 54                    let prompt_title = prompt.title?.to_string();
 55                    Some(ArgumentCompletion {
 56                        label: prompt_title.clone().into(),
 57                        new_text: prompt_title,
 58                        after_completion: true.into(),
 59                        replace_previous_arguments: true,
 60                    })
 61                })
 62                .collect())
 63        })
 64    }
 65
 66    fn run(
 67        self: Arc<Self>,
 68        arguments: &[String],
 69        _context_slash_command_output_sections: &[SlashCommandOutputSection<language::Anchor>],
 70        _context_buffer: BufferSnapshot,
 71        _workspace: WeakEntity<Workspace>,
 72        _delegate: Option<Arc<dyn LspAdapterDelegate>>,
 73        _: &mut Window,
 74        cx: &mut App,
 75    ) -> Task<SlashCommandResult> {
 76        let title = arguments.to_owned().join(" ");
 77        if title.trim().is_empty() {
 78            return Task::ready(Err(anyhow!("missing prompt name")));
 79        };
 80
 81        let store = PromptStore::global(cx);
 82        let title = SharedString::from(title.clone());
 83        let prompt = cx.spawn({
 84            let title = title.clone();
 85            async move |cx| {
 86                let store = store.await?;
 87                let body = store
 88                    .read_with(cx, |store, cx| {
 89                        let prompt_id = store
 90                            .id_for_title(&title)
 91                            .with_context(|| format!("no prompt found with title {:?}", title))?;
 92                        anyhow::Ok(store.load(prompt_id, cx))
 93                    })??
 94                    .await?;
 95                anyhow::Ok(body)
 96            }
 97        });
 98        cx.foreground_executor().spawn(async move {
 99            let mut prompt = prompt.await?;
100
101            if prompt.starts_with('/') {
102                // Prevent an edge case where the inserted prompt starts with a slash command (that leads to funky rendering).
103                prompt.insert(0, '\n');
104            }
105            if prompt.is_empty() {
106                prompt.push('\n');
107            }
108            let range = 0..prompt.len();
109            Ok(SlashCommandOutput {
110                text: prompt,
111                sections: vec![SlashCommandOutputSection {
112                    range,
113                    icon: IconName::Library,
114                    label: title,
115                    metadata: None,
116                }],
117                run_commands_in_text: true,
118            }
119            .to_event_stream())
120        })
121    }
122}