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