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