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