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