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::CompletionProvider;
 19pub use context::*;
 20pub use context_store::*;
 21use fs::Fs;
 22use gpui::{
 23    actions, impl_actions, AppContext, BorrowAppContext, Global, SharedString, UpdateGlobal,
 24};
 25use indexed_docs::IndexedDocsRegistry;
 26pub(crate) use inline_assistant::*;
 27use language_model::LanguageModelResponseMessage;
 28pub(crate) use model_selector::*;
 29use semantic_index::{CloudEmbeddingProvider, SemanticIndex};
 30use serde::{Deserialize, Serialize};
 31use settings::{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    cx.spawn(|mut cx| {
169        let client = client.clone();
170        async move {
171            let embedding_provider = CloudEmbeddingProvider::new(client.clone());
172            let semantic_index = SemanticIndex::new(
173                paths::embeddings_dir().join("semantic-index-db.0.mdb"),
174                Arc::new(embedding_provider),
175                &mut cx,
176            )
177            .await?;
178            cx.update(|cx| cx.set_global(semantic_index))
179        }
180    })
181    .detach();
182
183    context_store::init(&client);
184    prompt_library::init(cx);
185    init_completion_provider(Arc::clone(&client), cx);
186    assistant_slash_command::init(cx);
187    register_slash_commands(cx);
188    assistant_panel::init(cx);
189    inline_assistant::init(fs.clone(), client.telemetry().clone(), cx);
190    terminal_inline_assistant::init(fs.clone(), client.telemetry().clone(), cx);
191    IndexedDocsRegistry::init_global(cx);
192
193    CommandPaletteFilter::update_global(cx, |filter, _cx| {
194        filter.hide_namespace(Assistant::NAMESPACE);
195    });
196    Assistant::update_global(cx, |assistant, cx| {
197        let settings = AssistantSettings::get_global(cx);
198
199        assistant.set_enabled(settings.enabled, cx);
200    });
201    cx.observe_global::<SettingsStore>(|cx| {
202        Assistant::update_global(cx, |assistant, cx| {
203            let settings = AssistantSettings::get_global(cx);
204            assistant.set_enabled(settings.enabled, cx);
205        });
206    })
207    .detach();
208}
209
210fn init_completion_provider(client: Arc<Client>, cx: &mut AppContext) {
211    let provider = assistant_settings::create_provider_from_settings(client.clone(), 0, cx);
212    cx.set_global(CompletionProvider::new(provider, Some(client)));
213
214    let mut settings_version = 0;
215    cx.observe_global::<SettingsStore>(move |cx| {
216        settings_version += 1;
217        cx.update_global::<CompletionProvider, _>(|provider, cx| {
218            assistant_settings::update_completion_provider_settings(provider, settings_version, cx);
219        })
220    })
221    .detach();
222}
223
224fn register_slash_commands(cx: &mut AppContext) {
225    let slash_command_registry = SlashCommandRegistry::global(cx);
226    slash_command_registry.register_command(file_command::FileSlashCommand, true);
227    slash_command_registry.register_command(active_command::ActiveSlashCommand, true);
228    slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
229    slash_command_registry.register_command(tabs_command::TabsSlashCommand, true);
230    slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
231    slash_command_registry.register_command(search_command::SearchSlashCommand, true);
232    slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
233    slash_command_registry.register_command(default_command::DefaultSlashCommand, true);
234    slash_command_registry.register_command(term_command::TermSlashCommand, true);
235    slash_command_registry.register_command(now_command::NowSlashCommand, true);
236    slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
237    slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
238    slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
239}
240
241pub fn humanize_token_count(count: usize) -> String {
242    match count {
243        0..=999 => count.to_string(),
244        1000..=9999 => {
245            let thousands = count / 1000;
246            let hundreds = (count % 1000 + 50) / 100;
247            if hundreds == 0 {
248                format!("{}k", thousands)
249            } else if hundreds == 10 {
250                format!("{}k", thousands + 1)
251            } else {
252                format!("{}.{}k", thousands, hundreds)
253            }
254        }
255        _ => format!("{}k", (count + 500) / 1000),
256    }
257}
258
259#[cfg(test)]
260#[ctor::ctor]
261fn init_logger() {
262    if std::env::var("RUST_LOG").is_ok() {
263        env_logger::init();
264    }
265}