assistant.rs

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