assistant.rs

  1#![cfg_attr(target_os = "windows", allow(unused, dead_code))]
  2
  3pub mod assistant_panel;
  4pub mod assistant_settings;
  5mod context;
  6pub mod context_store;
  7mod inline_assistant;
  8mod model_selector;
  9mod prompt_library;
 10mod prompts;
 11mod slash_command;
 12pub(crate) mod slash_command_picker;
 13pub mod slash_command_settings;
 14mod streaming_diff;
 15mod terminal_inline_assistant;
 16mod tools;
 17mod workflow;
 18
 19pub use assistant_panel::{AssistantPanel, AssistantPanelEvent};
 20use assistant_settings::AssistantSettings;
 21use assistant_slash_command::SlashCommandRegistry;
 22use assistant_tool::ToolRegistry;
 23use client::{proto, Client};
 24use command_palette_hooks::CommandPaletteFilter;
 25pub use context::*;
 26use context_servers::ContextServerRegistry;
 27pub use context_store::*;
 28use feature_flags::FeatureFlagAppExt;
 29use fs::Fs;
 30use gpui::{actions, AppContext, Global, SharedString, UpdateGlobal};
 31use gpui::{impl_actions, Context as _};
 32use indexed_docs::IndexedDocsRegistry;
 33pub(crate) use inline_assistant::*;
 34use language_model::{
 35    LanguageModelId, LanguageModelProviderId, LanguageModelRegistry, LanguageModelResponseMessage,
 36};
 37pub(crate) use model_selector::*;
 38pub use prompts::PromptBuilder;
 39use prompts::PromptLoadingParams;
 40use semantic_index::{CloudEmbeddingProvider, SemanticDb};
 41use serde::{Deserialize, Serialize};
 42use settings::{update_settings_file, Settings, SettingsStore};
 43use slash_command::{
 44    auto_command, context_server_command, default_command, delta_command, diagnostics_command,
 45    docs_command, fetch_command, file_command, now_command, project_command, prompt_command,
 46    search_command, symbols_command, tab_command, terminal_command, workflow_command,
 47};
 48use std::path::PathBuf;
 49use std::sync::Arc;
 50pub(crate) use streaming_diff::*;
 51use util::ResultExt;
 52pub use workflow::*;
 53
 54use crate::slash_command_settings::SlashCommandSettings;
 55
 56actions!(
 57    assistant,
 58    [
 59        Assist,
 60        Split,
 61        CycleMessageRole,
 62        QuoteSelection,
 63        InsertIntoEditor,
 64        ToggleFocus,
 65        InsertActivePrompt,
 66        DeployHistory,
 67        DeployPromptLibrary,
 68        ConfirmCommand,
 69        NewContext,
 70        ToggleModelSelector,
 71    ]
 72);
 73
 74#[derive(PartialEq, Clone, Deserialize)]
 75pub enum InsertDraggedFiles {
 76    ProjectPaths(Vec<PathBuf>),
 77    ExternalFiles(Vec<PathBuf>),
 78}
 79
 80impl_actions!(assistant, [InsertDraggedFiles]);
 81
 82const DEFAULT_CONTEXT_LINES: usize = 50;
 83
 84#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 85pub struct MessageId(clock::Lamport);
 86
 87impl MessageId {
 88    pub fn as_u64(self) -> u64 {
 89        self.0.as_u64()
 90    }
 91}
 92
 93#[derive(Deserialize, Debug)]
 94pub struct LanguageModelUsage {
 95    pub prompt_tokens: u32,
 96    pub completion_tokens: u32,
 97    pub total_tokens: u32,
 98}
 99
100#[derive(Deserialize, Debug)]
101pub struct LanguageModelChoiceDelta {
102    pub index: u32,
103    pub delta: LanguageModelResponseMessage,
104    pub finish_reason: Option<String>,
105}
106
107#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
108pub enum MessageStatus {
109    Pending,
110    Done,
111    Error(SharedString),
112    Canceled,
113}
114
115impl MessageStatus {
116    pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
117        match status.variant {
118            Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
119            Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
120            Some(proto::context_message_status::Variant::Error(error)) => {
121                MessageStatus::Error(error.message.into())
122            }
123            Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
124            None => MessageStatus::Pending,
125        }
126    }
127
128    pub fn to_proto(&self) -> proto::ContextMessageStatus {
129        match self {
130            MessageStatus::Pending => proto::ContextMessageStatus {
131                variant: Some(proto::context_message_status::Variant::Pending(
132                    proto::context_message_status::Pending {},
133                )),
134            },
135            MessageStatus::Done => proto::ContextMessageStatus {
136                variant: Some(proto::context_message_status::Variant::Done(
137                    proto::context_message_status::Done {},
138                )),
139            },
140            MessageStatus::Error(message) => proto::ContextMessageStatus {
141                variant: Some(proto::context_message_status::Variant::Error(
142                    proto::context_message_status::Error {
143                        message: message.to_string(),
144                    },
145                )),
146            },
147            MessageStatus::Canceled => proto::ContextMessageStatus {
148                variant: Some(proto::context_message_status::Variant::Canceled(
149                    proto::context_message_status::Canceled {},
150                )),
151            },
152        }
153    }
154}
155
156/// The state pertaining to the Assistant.
157#[derive(Default)]
158struct Assistant {
159    /// Whether the Assistant is enabled.
160    enabled: bool,
161}
162
163impl Global for Assistant {}
164
165impl Assistant {
166    const NAMESPACE: &'static str = "assistant";
167
168    fn set_enabled(&mut self, enabled: bool, cx: &mut AppContext) {
169        if self.enabled == enabled {
170            return;
171        }
172
173        self.enabled = enabled;
174
175        if !enabled {
176            CommandPaletteFilter::update_global(cx, |filter, _cx| {
177                filter.hide_namespace(Self::NAMESPACE);
178            });
179
180            return;
181        }
182
183        CommandPaletteFilter::update_global(cx, |filter, _cx| {
184            filter.show_namespace(Self::NAMESPACE);
185        });
186    }
187}
188
189pub fn init(
190    fs: Arc<dyn Fs>,
191    client: Arc<Client>,
192    stdout_is_a_pty: bool,
193    cx: &mut AppContext,
194) -> Arc<PromptBuilder> {
195    cx.set_global(Assistant::default());
196    AssistantSettings::register(cx);
197    SlashCommandSettings::register(cx);
198
199    // TODO: remove this when 0.148.0 is released.
200    if AssistantSettings::get_global(cx).using_outdated_settings_version {
201        update_settings_file::<AssistantSettings>(fs.clone(), cx, {
202            let fs = fs.clone();
203            |content, cx| {
204                content.update_file(fs, cx);
205            }
206        });
207    }
208
209    cx.spawn(|mut cx| {
210        let client = client.clone();
211        async move {
212            let embedding_provider = CloudEmbeddingProvider::new(client.clone());
213            let semantic_index = SemanticDb::new(
214                paths::embeddings_dir().join("semantic-index-db.0.mdb"),
215                Arc::new(embedding_provider),
216                &mut cx,
217            )
218            .await?;
219
220            cx.update(|cx| cx.set_global(semantic_index))
221        }
222    })
223    .detach();
224
225    context_store::init(&client.clone().into());
226    prompt_library::init(cx);
227    init_language_model_settings(cx);
228    assistant_slash_command::init(cx);
229    assistant_tool::init(cx);
230    assistant_panel::init(cx);
231    context_servers::init(cx);
232
233    let prompt_builder = prompts::PromptBuilder::new(Some(PromptLoadingParams {
234        fs: fs.clone(),
235        repo_path: stdout_is_a_pty
236            .then(|| std::env::current_dir().log_err())
237            .flatten(),
238        cx,
239    }))
240    .log_err()
241    .map(Arc::new)
242    .unwrap_or_else(|| Arc::new(prompts::PromptBuilder::new(None).unwrap()));
243    register_slash_commands(Some(prompt_builder.clone()), cx);
244    register_tools(cx);
245    inline_assistant::init(
246        fs.clone(),
247        prompt_builder.clone(),
248        client.telemetry().clone(),
249        cx,
250    );
251    terminal_inline_assistant::init(
252        fs.clone(),
253        prompt_builder.clone(),
254        client.telemetry().clone(),
255        cx,
256    );
257    IndexedDocsRegistry::init_global(cx);
258
259    CommandPaletteFilter::update_global(cx, |filter, _cx| {
260        filter.hide_namespace(Assistant::NAMESPACE);
261    });
262    Assistant::update_global(cx, |assistant, cx| {
263        let settings = AssistantSettings::get_global(cx);
264
265        assistant.set_enabled(settings.enabled, cx);
266    });
267    cx.observe_global::<SettingsStore>(|cx| {
268        Assistant::update_global(cx, |assistant, cx| {
269            let settings = AssistantSettings::get_global(cx);
270            assistant.set_enabled(settings.enabled, cx);
271        });
272    })
273    .detach();
274
275    register_context_server_handlers(cx);
276
277    prompt_builder
278}
279
280fn register_context_server_handlers(cx: &mut AppContext) {
281    cx.subscribe(
282        &context_servers::manager::ContextServerManager::global(cx),
283        |manager, event, cx| match event {
284            context_servers::manager::Event::ServerStarted { server_id } => {
285                cx.update_model(
286                    &manager,
287                    |manager: &mut context_servers::manager::ContextServerManager, cx| {
288                        let slash_command_registry = SlashCommandRegistry::global(cx);
289                        let context_server_registry = ContextServerRegistry::global(cx);
290                        if let Some(server) = manager.get_server(server_id) {
291                            cx.spawn(|_, _| async move {
292                                let Some(protocol) = server.client.read().clone() else {
293                                    return;
294                                };
295
296                                if let Some(prompts) = protocol.list_prompts().await.log_err() {
297                                    for prompt in prompts
298                                        .into_iter()
299                                        .filter(context_server_command::acceptable_prompt)
300                                    {
301                                        log::info!(
302                                            "registering context server command: {:?}",
303                                            prompt.name
304                                        );
305                                        context_server_registry.register_command(
306                                            server.id.clone(),
307                                            prompt.name.as_str(),
308                                        );
309                                        slash_command_registry.register_command(
310                                            context_server_command::ContextServerSlashCommand::new(
311                                                &server, prompt,
312                                            ),
313                                            true,
314                                        );
315                                    }
316                                }
317                            })
318                            .detach();
319                        }
320                    },
321                );
322            }
323            context_servers::manager::Event::ServerStopped { server_id } => {
324                let slash_command_registry = SlashCommandRegistry::global(cx);
325                let context_server_registry = ContextServerRegistry::global(cx);
326                if let Some(commands) = context_server_registry.get_commands(server_id) {
327                    for command_name in commands {
328                        slash_command_registry.unregister_command_by_name(&command_name);
329                        context_server_registry.unregister_command(&server_id, &command_name);
330                    }
331                }
332            }
333        },
334    )
335    .detach();
336}
337
338fn init_language_model_settings(cx: &mut AppContext) {
339    update_active_language_model_from_settings(cx);
340
341    cx.observe_global::<SettingsStore>(update_active_language_model_from_settings)
342        .detach();
343    cx.subscribe(
344        &LanguageModelRegistry::global(cx),
345        |_, event: &language_model::Event, cx| match event {
346            language_model::Event::ProviderStateChanged
347            | language_model::Event::AddedProvider(_)
348            | language_model::Event::RemovedProvider(_) => {
349                update_active_language_model_from_settings(cx);
350            }
351            _ => {}
352        },
353    )
354    .detach();
355}
356
357fn update_active_language_model_from_settings(cx: &mut AppContext) {
358    let settings = AssistantSettings::get_global(cx);
359    let provider_name = LanguageModelProviderId::from(settings.default_model.provider.clone());
360    let model_id = LanguageModelId::from(settings.default_model.model.clone());
361    LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
362        registry.select_active_model(&provider_name, &model_id, cx);
363    });
364}
365
366fn register_slash_commands(prompt_builder: Option<Arc<PromptBuilder>>, cx: &mut AppContext) {
367    let slash_command_registry = SlashCommandRegistry::global(cx);
368
369    slash_command_registry.register_command(file_command::FileSlashCommand, true);
370    slash_command_registry.register_command(delta_command::DeltaSlashCommand, true);
371    slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
372    slash_command_registry.register_command(tab_command::TabSlashCommand, true);
373    slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
374    slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
375    slash_command_registry.register_command(default_command::DefaultSlashCommand, false);
376    slash_command_registry.register_command(terminal_command::TerminalSlashCommand, true);
377    slash_command_registry.register_command(now_command::NowSlashCommand, false);
378    slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
379
380    if let Some(prompt_builder) = prompt_builder {
381        slash_command_registry.register_command(
382            workflow_command::WorkflowSlashCommand::new(prompt_builder.clone()),
383            true,
384        );
385    }
386    slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
387
388    cx.observe_flag::<auto_command::AutoSlashCommandFeatureFlag, _>({
389        let slash_command_registry = slash_command_registry.clone();
390        move |is_enabled, _cx| {
391            if is_enabled {
392                // [#auto-staff-ship] TODO remove this when /auto is no longer staff-shipped
393                slash_command_registry.register_command(auto_command::AutoCommand, true);
394            }
395        }
396    })
397    .detach();
398
399    update_slash_commands_from_settings(cx);
400    cx.observe_global::<SettingsStore>(update_slash_commands_from_settings)
401        .detach();
402
403    cx.observe_flag::<search_command::SearchSlashCommandFeatureFlag, _>({
404        let slash_command_registry = slash_command_registry.clone();
405        move |is_enabled, _cx| {
406            if is_enabled {
407                slash_command_registry.register_command(search_command::SearchSlashCommand, true);
408            }
409        }
410    })
411    .detach();
412}
413
414fn update_slash_commands_from_settings(cx: &mut AppContext) {
415    let slash_command_registry = SlashCommandRegistry::global(cx);
416    let settings = SlashCommandSettings::get_global(cx);
417
418    if settings.docs.enabled {
419        slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
420    } else {
421        slash_command_registry.unregister_command(docs_command::DocsSlashCommand);
422    }
423
424    if settings.project.enabled {
425        slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
426    } else {
427        slash_command_registry.unregister_command(project_command::ProjectSlashCommand);
428    }
429}
430
431fn register_tools(cx: &mut AppContext) {
432    let tool_registry = ToolRegistry::global(cx);
433    tool_registry.register_tool(tools::now_tool::NowTool);
434}
435
436pub fn humanize_token_count(count: usize) -> String {
437    match count {
438        0..=999 => count.to_string(),
439        1000..=9999 => {
440            let thousands = count / 1000;
441            let hundreds = (count % 1000 + 50) / 100;
442            if hundreds == 0 {
443                format!("{}k", thousands)
444            } else if hundreds == 10 {
445                format!("{}k", thousands + 1)
446            } else {
447                format!("{}.{}k", thousands, hundreds)
448            }
449        }
450        _ => format!("{}k", (count + 500) / 1000),
451    }
452}
453
454#[cfg(test)]
455#[ctor::ctor]
456fn init_logger() {
457    if std::env::var("RUST_LOG").is_ok() {
458        env_logger::init();
459    }
460}