assistant.rs

  1pub mod assistant_panel;
  2pub mod assistant_settings;
  3mod context;
  4pub mod context_store;
  5mod inline_assistant;
  6mod model_selector;
  7mod prompt_library;
  8mod prompts;
  9mod slash_command;
 10mod streaming_diff;
 11mod terminal_inline_assistant;
 12
 13pub use assistant_panel::{AssistantPanel, AssistantPanelEvent};
 14use assistant_settings::AssistantSettings;
 15use assistant_slash_command::SlashCommandRegistry;
 16use client::{proto, Client};
 17use command_palette_hooks::CommandPaletteFilter;
 18use completion::LanguageModelCompletionProvider;
 19pub use context::*;
 20pub use context_store::*;
 21use fs::Fs;
 22use gpui::{actions, impl_actions, AppContext, Global, SharedString, UpdateGlobal};
 23use indexed_docs::IndexedDocsRegistry;
 24pub(crate) use inline_assistant::*;
 25use language_model::{
 26    LanguageModelId, LanguageModelProviderId, LanguageModelRegistry, LanguageModelResponseMessage,
 27};
 28pub(crate) use model_selector::*;
 29use semantic_index::{CloudEmbeddingProvider, SemanticIndex};
 30use serde::{Deserialize, Serialize};
 31use settings::{update_settings_file, Settings, SettingsStore};
 32use slash_command::{
 33    active_command, default_command, diagnostics_command, docs_command, fetch_command,
 34    file_command, now_command, project_command, prompt_command, search_command, symbols_command,
 35    tabs_command, term_command,
 36};
 37use std::sync::Arc;
 38pub(crate) use streaming_diff::*;
 39
 40actions!(
 41    assistant,
 42    [
 43        Assist,
 44        Split,
 45        CycleMessageRole,
 46        QuoteSelection,
 47        InsertIntoEditor,
 48        ToggleFocus,
 49        ResetKey,
 50        InsertActivePrompt,
 51        DeployHistory,
 52        DeployPromptLibrary,
 53        ConfirmCommand,
 54        ToggleModelSelector,
 55        DebugEditSteps
 56    ]
 57);
 58
 59#[derive(Clone, Default, Deserialize, PartialEq)]
 60pub struct InlineAssist {
 61    prompt: Option<String>,
 62}
 63
 64impl_actions!(assistant, [InlineAssist]);
 65
 66#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
 67pub struct MessageId(clock::Lamport);
 68
 69impl MessageId {
 70    pub fn as_u64(self) -> u64 {
 71        self.0.as_u64()
 72    }
 73}
 74
 75#[derive(Deserialize, Debug)]
 76pub struct LanguageModelUsage {
 77    pub prompt_tokens: u32,
 78    pub completion_tokens: u32,
 79    pub total_tokens: u32,
 80}
 81
 82#[derive(Deserialize, Debug)]
 83pub struct LanguageModelChoiceDelta {
 84    pub index: u32,
 85    pub delta: LanguageModelResponseMessage,
 86    pub finish_reason: Option<String>,
 87}
 88
 89#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 90pub enum MessageStatus {
 91    Pending,
 92    Done,
 93    Error(SharedString),
 94}
 95
 96impl MessageStatus {
 97    pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
 98        match status.variant {
 99            Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
100            Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
101            Some(proto::context_message_status::Variant::Error(error)) => {
102                MessageStatus::Error(error.message.into())
103            }
104            None => MessageStatus::Pending,
105        }
106    }
107
108    pub fn to_proto(&self) -> proto::ContextMessageStatus {
109        match self {
110            MessageStatus::Pending => proto::ContextMessageStatus {
111                variant: Some(proto::context_message_status::Variant::Pending(
112                    proto::context_message_status::Pending {},
113                )),
114            },
115            MessageStatus::Done => proto::ContextMessageStatus {
116                variant: Some(proto::context_message_status::Variant::Done(
117                    proto::context_message_status::Done {},
118                )),
119            },
120            MessageStatus::Error(message) => proto::ContextMessageStatus {
121                variant: Some(proto::context_message_status::Variant::Error(
122                    proto::context_message_status::Error {
123                        message: message.to_string(),
124                    },
125                )),
126            },
127        }
128    }
129}
130
131/// The state pertaining to the Assistant.
132#[derive(Default)]
133struct Assistant {
134    /// Whether the Assistant is enabled.
135    enabled: bool,
136}
137
138impl Global for Assistant {}
139
140impl Assistant {
141    const NAMESPACE: &'static str = "assistant";
142
143    fn set_enabled(&mut self, enabled: bool, cx: &mut AppContext) {
144        if self.enabled == enabled {
145            return;
146        }
147
148        self.enabled = enabled;
149
150        if !enabled {
151            CommandPaletteFilter::update_global(cx, |filter, _cx| {
152                filter.hide_namespace(Self::NAMESPACE);
153            });
154
155            return;
156        }
157
158        CommandPaletteFilter::update_global(cx, |filter, _cx| {
159            filter.show_namespace(Self::NAMESPACE);
160        });
161    }
162}
163
164pub fn init(fs: Arc<dyn Fs>, client: Arc<Client>, cx: &mut AppContext) {
165    cx.set_global(Assistant::default());
166    AssistantSettings::register(cx);
167
168    // TODO: remove this when 0.148.0 is released.
169    if AssistantSettings::get_global(cx).using_outdated_settings_version {
170        update_settings_file::<AssistantSettings>(fs.clone(), cx, {
171            let fs = fs.clone();
172            |content, cx| {
173                content.update_file(fs, cx);
174            }
175        });
176    }
177
178    cx.spawn(|mut cx| {
179        let client = client.clone();
180        async move {
181            let embedding_provider = CloudEmbeddingProvider::new(client.clone());
182            let semantic_index = SemanticIndex::new(
183                paths::embeddings_dir().join("semantic-index-db.0.mdb"),
184                Arc::new(embedding_provider),
185                &mut cx,
186            )
187            .await?;
188            cx.update(|cx| cx.set_global(semantic_index))
189        }
190    })
191    .detach();
192
193    context_store::init(&client);
194    prompt_library::init(cx);
195    init_completion_provider(cx);
196    assistant_slash_command::init(cx);
197    register_slash_commands(cx);
198    assistant_panel::init(cx);
199    inline_assistant::init(fs.clone(), client.telemetry().clone(), cx);
200    terminal_inline_assistant::init(fs.clone(), client.telemetry().clone(), cx);
201    IndexedDocsRegistry::init_global(cx);
202
203    CommandPaletteFilter::update_global(cx, |filter, _cx| {
204        filter.hide_namespace(Assistant::NAMESPACE);
205    });
206    Assistant::update_global(cx, |assistant, cx| {
207        let settings = AssistantSettings::get_global(cx);
208
209        assistant.set_enabled(settings.enabled, cx);
210    });
211    cx.observe_global::<SettingsStore>(|cx| {
212        Assistant::update_global(cx, |assistant, cx| {
213            let settings = AssistantSettings::get_global(cx);
214            assistant.set_enabled(settings.enabled, cx);
215        });
216    })
217    .detach();
218}
219
220fn init_completion_provider(cx: &mut AppContext) {
221    completion::init(cx);
222    update_active_language_model_from_settings(cx);
223
224    cx.observe_global::<SettingsStore>(update_active_language_model_from_settings)
225        .detach();
226    cx.observe(&LanguageModelRegistry::global(cx), |_, cx| {
227        update_active_language_model_from_settings(cx)
228    })
229    .detach();
230}
231
232fn update_active_language_model_from_settings(cx: &mut AppContext) {
233    let settings = AssistantSettings::get_global(cx);
234    let provider_name = LanguageModelProviderId::from(settings.default_model.provider.clone());
235    let model_id = LanguageModelId::from(settings.default_model.model.clone());
236
237    let Some(provider) = LanguageModelRegistry::global(cx)
238        .read(cx)
239        .provider(&provider_name)
240    else {
241        return;
242    };
243
244    let models = provider.provided_models(cx);
245    if let Some(model) = models.iter().find(|model| model.id() == model_id).cloned() {
246        LanguageModelCompletionProvider::global(cx).update(cx, |completion_provider, cx| {
247            completion_provider.set_active_model(model, cx);
248        });
249    }
250}
251
252fn register_slash_commands(cx: &mut AppContext) {
253    let slash_command_registry = SlashCommandRegistry::global(cx);
254    slash_command_registry.register_command(file_command::FileSlashCommand, true);
255    slash_command_registry.register_command(active_command::ActiveSlashCommand, true);
256    slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
257    slash_command_registry.register_command(tabs_command::TabsSlashCommand, true);
258    slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
259    slash_command_registry.register_command(search_command::SearchSlashCommand, true);
260    slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
261    slash_command_registry.register_command(default_command::DefaultSlashCommand, true);
262    slash_command_registry.register_command(term_command::TermSlashCommand, true);
263    slash_command_registry.register_command(now_command::NowSlashCommand, true);
264    slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
265    slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
266    slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
267}
268
269pub fn humanize_token_count(count: usize) -> String {
270    match count {
271        0..=999 => count.to_string(),
272        1000..=9999 => {
273            let thousands = count / 1000;
274            let hundreds = (count % 1000 + 50) / 100;
275            if hundreds == 0 {
276                format!("{}k", thousands)
277            } else if hundreds == 10 {
278                format!("{}k", thousands + 1)
279            } else {
280                format!("{}.{}k", thousands, hundreds)
281            }
282        }
283        _ => format!("{}k", (count + 500) / 1000),
284    }
285}
286
287#[cfg(test)]
288#[ctor::ctor]
289fn init_logger() {
290    if std::env::var("RUST_LOG").is_ok() {
291        env_logger::init();
292    }
293}