assistant.rs

  1#![cfg_attr(target_os = "windows", allow(unused, dead_code))]
  2
  3pub mod assistant_panel;
  4pub mod assistant_settings;
  5mod context;
  6pub(crate) mod context_inspector;
  7pub mod context_store;
  8mod inline_assistant;
  9mod model_selector;
 10mod prompt_library;
 11mod prompts;
 12mod slash_command;
 13pub mod slash_command_settings;
 14mod streaming_diff;
 15mod terminal_inline_assistant;
 16mod workflow;
 17
 18pub use assistant_panel::{AssistantPanel, AssistantPanelEvent};
 19use assistant_settings::AssistantSettings;
 20use assistant_slash_command::SlashCommandRegistry;
 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::{actions, impl_actions, AppContext, Global, SharedString, UpdateGlobal};
 28use indexed_docs::IndexedDocsRegistry;
 29pub(crate) use inline_assistant::*;
 30use language_model::{
 31    LanguageModelId, LanguageModelProviderId, LanguageModelRegistry, LanguageModelResponseMessage,
 32};
 33pub(crate) use model_selector::*;
 34pub use prompts::PromptBuilder;
 35use prompts::PromptOverrideContext;
 36use semantic_index::{CloudEmbeddingProvider, SemanticIndex};
 37use serde::{Deserialize, Serialize};
 38use settings::{update_settings_file, Settings, SettingsStore};
 39use slash_command::{
 40    default_command, diagnostics_command, docs_command, fetch_command, file_command, now_command,
 41    project_command, prompt_command, search_command, symbols_command, tab_command,
 42    terminal_command, workflow_command,
 43};
 44use std::sync::Arc;
 45pub(crate) use streaming_diff::*;
 46use util::ResultExt;
 47pub use workflow::*;
 48
 49use crate::slash_command_settings::SlashCommandSettings;
 50
 51actions!(
 52    assistant,
 53    [
 54        Assist,
 55        Split,
 56        CycleMessageRole,
 57        QuoteSelection,
 58        InsertIntoEditor,
 59        ToggleFocus,
 60        InsertActivePrompt,
 61        ShowConfiguration,
 62        DeployHistory,
 63        DeployPromptLibrary,
 64        ConfirmCommand,
 65        ToggleModelSelector,
 66        DebugWorkflowSteps
 67    ]
 68);
 69
 70const DEFAULT_CONTEXT_LINES: usize = 50;
 71
 72#[derive(Clone, Default, Deserialize, PartialEq)]
 73pub struct InlineAssist {
 74    prompt: Option<String>,
 75}
 76
 77impl_actions!(assistant, [InlineAssist]);
 78
 79#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 80pub struct MessageId(clock::Lamport);
 81
 82impl MessageId {
 83    pub fn as_u64(self) -> u64 {
 84        self.0.as_u64()
 85    }
 86}
 87
 88#[derive(Deserialize, Debug)]
 89pub struct LanguageModelUsage {
 90    pub prompt_tokens: u32,
 91    pub completion_tokens: u32,
 92    pub total_tokens: u32,
 93}
 94
 95#[derive(Deserialize, Debug)]
 96pub struct LanguageModelChoiceDelta {
 97    pub index: u32,
 98    pub delta: LanguageModelResponseMessage,
 99    pub finish_reason: Option<String>,
100}
101
102#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
103pub enum MessageStatus {
104    Pending,
105    Done,
106    Error(SharedString),
107    Canceled,
108}
109
110impl MessageStatus {
111    pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
112        match status.variant {
113            Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
114            Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
115            Some(proto::context_message_status::Variant::Error(error)) => {
116                MessageStatus::Error(error.message.into())
117            }
118            Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
119            None => MessageStatus::Pending,
120        }
121    }
122
123    pub fn to_proto(&self) -> proto::ContextMessageStatus {
124        match self {
125            MessageStatus::Pending => proto::ContextMessageStatus {
126                variant: Some(proto::context_message_status::Variant::Pending(
127                    proto::context_message_status::Pending {},
128                )),
129            },
130            MessageStatus::Done => proto::ContextMessageStatus {
131                variant: Some(proto::context_message_status::Variant::Done(
132                    proto::context_message_status::Done {},
133                )),
134            },
135            MessageStatus::Error(message) => proto::ContextMessageStatus {
136                variant: Some(proto::context_message_status::Variant::Error(
137                    proto::context_message_status::Error {
138                        message: message.to_string(),
139                    },
140                )),
141            },
142            MessageStatus::Canceled => proto::ContextMessageStatus {
143                variant: Some(proto::context_message_status::Variant::Canceled(
144                    proto::context_message_status::Canceled {},
145                )),
146            },
147        }
148    }
149}
150
151/// The state pertaining to the Assistant.
152#[derive(Default)]
153struct Assistant {
154    /// Whether the Assistant is enabled.
155    enabled: bool,
156}
157
158impl Global for Assistant {}
159
160impl Assistant {
161    const NAMESPACE: &'static str = "assistant";
162
163    fn set_enabled(&mut self, enabled: bool, cx: &mut AppContext) {
164        if self.enabled == enabled {
165            return;
166        }
167
168        self.enabled = enabled;
169
170        if !enabled {
171            CommandPaletteFilter::update_global(cx, |filter, _cx| {
172                filter.hide_namespace(Self::NAMESPACE);
173            });
174
175            return;
176        }
177
178        CommandPaletteFilter::update_global(cx, |filter, _cx| {
179            filter.show_namespace(Self::NAMESPACE);
180        });
181    }
182}
183
184pub fn init(
185    fs: Arc<dyn Fs>,
186    client: Arc<Client>,
187    dev_mode: bool,
188    cx: &mut AppContext,
189) -> Arc<PromptBuilder> {
190    cx.set_global(Assistant::default());
191    AssistantSettings::register(cx);
192    SlashCommandSettings::register(cx);
193
194    // TODO: remove this when 0.148.0 is released.
195    if AssistantSettings::get_global(cx).using_outdated_settings_version {
196        update_settings_file::<AssistantSettings>(fs.clone(), cx, {
197            let fs = fs.clone();
198            |content, cx| {
199                content.update_file(fs, cx);
200            }
201        });
202    }
203
204    cx.spawn(|mut cx| {
205        let client = client.clone();
206        async move {
207            let embedding_provider = CloudEmbeddingProvider::new(client.clone());
208            let semantic_index = SemanticIndex::new(
209                paths::embeddings_dir().join("semantic-index-db.0.mdb"),
210                Arc::new(embedding_provider),
211                &mut cx,
212            )
213            .await?;
214            cx.update(|cx| cx.set_global(semantic_index))
215        }
216    })
217    .detach();
218
219    context_store::init(&client);
220    prompt_library::init(cx);
221    init_language_model_settings(cx);
222    assistant_slash_command::init(cx);
223    assistant_panel::init(cx);
224
225    let prompt_builder = prompts::PromptBuilder::new(Some(PromptOverrideContext {
226        dev_mode,
227        fs: fs.clone(),
228        cx,
229    }))
230    .log_err()
231    .map(Arc::new)
232    .unwrap_or_else(|| Arc::new(prompts::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    IndexedDocsRegistry::init_global(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    LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
291        registry.select_active_model(&provider_name, &model_id, cx);
292    });
293}
294
295fn register_slash_commands(prompt_builder: Option<Arc<PromptBuilder>>, cx: &mut AppContext) {
296    let slash_command_registry = SlashCommandRegistry::global(cx);
297    slash_command_registry.register_command(file_command::FileSlashCommand, true);
298    slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
299    slash_command_registry.register_command(tab_command::TabSlashCommand, true);
300    slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
301    slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
302    slash_command_registry.register_command(default_command::DefaultSlashCommand, false);
303    slash_command_registry.register_command(terminal_command::TerminalSlashCommand, true);
304    slash_command_registry.register_command(now_command::NowSlashCommand, false);
305    slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
306
307    if let Some(prompt_builder) = prompt_builder {
308        slash_command_registry.register_command(
309            workflow_command::WorkflowSlashCommand::new(prompt_builder),
310            true,
311        );
312    }
313    slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
314
315    update_slash_commands_from_settings(cx);
316    cx.observe_global::<SettingsStore>(update_slash_commands_from_settings)
317        .detach();
318
319    cx.observe_flag::<search_command::SearchSlashCommandFeatureFlag, _>({
320        let slash_command_registry = slash_command_registry.clone();
321        move |is_enabled, _cx| {
322            if is_enabled {
323                slash_command_registry.register_command(search_command::SearchSlashCommand, true);
324            }
325        }
326    })
327    .detach();
328}
329
330fn update_slash_commands_from_settings(cx: &mut AppContext) {
331    let slash_command_registry = SlashCommandRegistry::global(cx);
332    let settings = SlashCommandSettings::get_global(cx);
333
334    if settings.docs.enabled {
335        slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
336    } else {
337        slash_command_registry.unregister_command(docs_command::DocsSlashCommand);
338    }
339
340    if settings.project.enabled {
341        slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
342    } else {
343        slash_command_registry.unregister_command(project_command::ProjectSlashCommand);
344    }
345}
346
347pub fn humanize_token_count(count: usize) -> String {
348    match count {
349        0..=999 => count.to_string(),
350        1000..=9999 => {
351            let thousands = count / 1000;
352            let hundreds = (count % 1000 + 50) / 100;
353            if hundreds == 0 {
354                format!("{}k", thousands)
355            } else if hundreds == 10 {
356                format!("{}k", thousands + 1)
357            } else {
358                format!("{}.{}k", thousands, hundreds)
359            }
360        }
361        _ => format!("{}k", (count + 500) / 1000),
362    }
363}
364
365#[cfg(test)]
366#[ctor::ctor]
367fn init_logger() {
368    if std::env::var("RUST_LOG").is_ok() {
369        env_logger::init();
370    }
371}