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 model_selector;
9mod prompt_library;
10mod prompts;
11mod slash_command;
12pub(crate) mod slash_command_picker;
13pub mod slash_command_settings;
14mod streaming_diff;
15mod terminal_inline_assistant;
16mod tools;
17mod workflow;
18
19pub use assistant_panel::{AssistantPanel, AssistantPanelEvent};
20use assistant_settings::AssistantSettings;
21use assistant_slash_command::SlashCommandRegistry;
22use assistant_tool::ToolRegistry;
23use client::{proto, Client};
24use command_palette_hooks::CommandPaletteFilter;
25pub use context::*;
26use context_servers::ContextServerRegistry;
27pub use context_store::*;
28use feature_flags::FeatureFlagAppExt;
29use fs::Fs;
30use gpui::{actions, AppContext, Global, SharedString, UpdateGlobal};
31use gpui::{impl_actions, Context as _};
32use indexed_docs::IndexedDocsRegistry;
33pub(crate) use inline_assistant::*;
34use language_model::{
35 LanguageModelId, LanguageModelProviderId, LanguageModelRegistry, LanguageModelResponseMessage,
36};
37pub(crate) use model_selector::*;
38pub use prompts::PromptBuilder;
39use prompts::PromptLoadingParams;
40use semantic_index::{CloudEmbeddingProvider, SemanticDb};
41use serde::{Deserialize, Serialize};
42use settings::{update_settings_file, Settings, SettingsStore};
43use slash_command::{
44 auto_command, context_server_command, default_command, delta_command, diagnostics_command,
45 docs_command, fetch_command, file_command, now_command, project_command, prompt_command,
46 search_command, symbols_command, tab_command, terminal_command, workflow_command,
47};
48use std::path::PathBuf;
49use std::sync::Arc;
50pub(crate) use streaming_diff::*;
51use util::ResultExt;
52pub use workflow::*;
53
54use crate::slash_command_settings::SlashCommandSettings;
55
56actions!(
57 assistant,
58 [
59 Assist,
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 embedding_provider = CloudEmbeddingProvider::new(client.clone());
216 let semantic_index = SemanticDb::new(
217 paths::embeddings_dir().join("semantic-index-db.0.mdb"),
218 Arc::new(embedding_provider),
219 &mut cx,
220 )
221 .await?;
222
223 cx.update(|cx| cx.set_global(semantic_index))
224 }
225 })
226 .detach();
227
228 context_store::init(&client.clone().into());
229 prompt_library::init(cx);
230 init_language_model_settings(cx);
231 assistant_slash_command::init(cx);
232 assistant_tool::init(cx);
233 assistant_panel::init(cx);
234 context_servers::init(cx);
235
236 let prompt_builder = prompts::PromptBuilder::new(Some(PromptLoadingParams {
237 fs: fs.clone(),
238 repo_path: stdout_is_a_pty
239 .then(|| std::env::current_dir().log_err())
240 .flatten(),
241 cx,
242 }))
243 .log_err()
244 .map(Arc::new)
245 .unwrap_or_else(|| Arc::new(prompts::PromptBuilder::new(None).unwrap()));
246 register_slash_commands(Some(prompt_builder.clone()), cx);
247 register_tools(cx);
248 inline_assistant::init(
249 fs.clone(),
250 prompt_builder.clone(),
251 client.telemetry().clone(),
252 cx,
253 );
254 terminal_inline_assistant::init(
255 fs.clone(),
256 prompt_builder.clone(),
257 client.telemetry().clone(),
258 cx,
259 );
260 IndexedDocsRegistry::init_global(cx);
261
262 CommandPaletteFilter::update_global(cx, |filter, _cx| {
263 filter.hide_namespace(Assistant::NAMESPACE);
264 });
265 Assistant::update_global(cx, |assistant, cx| {
266 let settings = AssistantSettings::get_global(cx);
267
268 assistant.set_enabled(settings.enabled, cx);
269 });
270 cx.observe_global::<SettingsStore>(|cx| {
271 Assistant::update_global(cx, |assistant, cx| {
272 let settings = AssistantSettings::get_global(cx);
273 assistant.set_enabled(settings.enabled, cx);
274 });
275 })
276 .detach();
277
278 register_context_server_handlers(cx);
279
280 prompt_builder
281}
282
283fn register_context_server_handlers(cx: &mut AppContext) {
284 cx.subscribe(
285 &context_servers::manager::ContextServerManager::global(cx),
286 |manager, event, cx| match event {
287 context_servers::manager::Event::ServerStarted { server_id } => {
288 cx.update_model(
289 &manager,
290 |manager: &mut context_servers::manager::ContextServerManager, cx| {
291 let slash_command_registry = SlashCommandRegistry::global(cx);
292 let context_server_registry = ContextServerRegistry::global(cx);
293 if let Some(server) = manager.get_server(server_id) {
294 cx.spawn(|_, _| async move {
295 let Some(protocol) = server.client.read().clone() else {
296 return;
297 };
298
299 if let Some(prompts) = protocol.list_prompts().await.log_err() {
300 for prompt in prompts
301 .into_iter()
302 .filter(context_server_command::acceptable_prompt)
303 {
304 log::info!(
305 "registering context server command: {:?}",
306 prompt.name
307 );
308 context_server_registry.register_command(
309 server.id.clone(),
310 prompt.name.as_str(),
311 );
312 slash_command_registry.register_command(
313 context_server_command::ContextServerSlashCommand::new(
314 &server, prompt,
315 ),
316 true,
317 );
318 }
319 }
320 })
321 .detach();
322 }
323 },
324 );
325 }
326 context_servers::manager::Event::ServerStopped { server_id } => {
327 let slash_command_registry = SlashCommandRegistry::global(cx);
328 let context_server_registry = ContextServerRegistry::global(cx);
329 if let Some(commands) = context_server_registry.get_commands(server_id) {
330 for command_name in commands {
331 slash_command_registry.unregister_command_by_name(&command_name);
332 context_server_registry.unregister_command(&server_id, &command_name);
333 }
334 }
335 }
336 },
337 )
338 .detach();
339}
340
341fn init_language_model_settings(cx: &mut AppContext) {
342 update_active_language_model_from_settings(cx);
343
344 cx.observe_global::<SettingsStore>(update_active_language_model_from_settings)
345 .detach();
346 cx.subscribe(
347 &LanguageModelRegistry::global(cx),
348 |_, event: &language_model::Event, cx| match event {
349 language_model::Event::ProviderStateChanged
350 | language_model::Event::AddedProvider(_)
351 | language_model::Event::RemovedProvider(_) => {
352 update_active_language_model_from_settings(cx);
353 }
354 _ => {}
355 },
356 )
357 .detach();
358}
359
360fn update_active_language_model_from_settings(cx: &mut AppContext) {
361 let settings = AssistantSettings::get_global(cx);
362 let provider_name = LanguageModelProviderId::from(settings.default_model.provider.clone());
363 let model_id = LanguageModelId::from(settings.default_model.model.clone());
364 let inline_alternatives = settings
365 .inline_alternatives
366 .iter()
367 .map(|alternative| {
368 (
369 LanguageModelProviderId::from(alternative.provider.clone()),
370 LanguageModelId::from(alternative.model.clone()),
371 )
372 })
373 .collect::<Vec<_>>();
374 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
375 registry.select_active_model(&provider_name, &model_id, cx);
376 registry.select_inline_alternative_models(inline_alternatives, cx);
377 });
378}
379
380fn register_slash_commands(prompt_builder: Option<Arc<PromptBuilder>>, cx: &mut AppContext) {
381 let slash_command_registry = SlashCommandRegistry::global(cx);
382
383 slash_command_registry.register_command(file_command::FileSlashCommand, true);
384 slash_command_registry.register_command(delta_command::DeltaSlashCommand, true);
385 slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
386 slash_command_registry.register_command(tab_command::TabSlashCommand, true);
387 slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
388 slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
389 slash_command_registry.register_command(default_command::DefaultSlashCommand, false);
390 slash_command_registry.register_command(terminal_command::TerminalSlashCommand, true);
391 slash_command_registry.register_command(now_command::NowSlashCommand, false);
392 slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
393
394 if let Some(prompt_builder) = prompt_builder {
395 slash_command_registry.register_command(
396 workflow_command::WorkflowSlashCommand::new(prompt_builder.clone()),
397 true,
398 );
399 }
400 slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
401
402 cx.observe_flag::<auto_command::AutoSlashCommandFeatureFlag, _>({
403 let slash_command_registry = slash_command_registry.clone();
404 move |is_enabled, _cx| {
405 if is_enabled {
406 // [#auto-staff-ship] TODO remove this when /auto is no longer staff-shipped
407 slash_command_registry.register_command(auto_command::AutoCommand, true);
408 }
409 }
410 })
411 .detach();
412
413 update_slash_commands_from_settings(cx);
414 cx.observe_global::<SettingsStore>(update_slash_commands_from_settings)
415 .detach();
416
417 cx.observe_flag::<search_command::SearchSlashCommandFeatureFlag, _>({
418 let slash_command_registry = slash_command_registry.clone();
419 move |is_enabled, _cx| {
420 if is_enabled {
421 slash_command_registry.register_command(search_command::SearchSlashCommand, true);
422 }
423 }
424 })
425 .detach();
426}
427
428fn update_slash_commands_from_settings(cx: &mut AppContext) {
429 let slash_command_registry = SlashCommandRegistry::global(cx);
430 let settings = SlashCommandSettings::get_global(cx);
431
432 if settings.docs.enabled {
433 slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
434 } else {
435 slash_command_registry.unregister_command(docs_command::DocsSlashCommand);
436 }
437
438 if settings.project.enabled {
439 slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
440 } else {
441 slash_command_registry.unregister_command(project_command::ProjectSlashCommand);
442 }
443}
444
445fn register_tools(cx: &mut AppContext) {
446 let tool_registry = ToolRegistry::global(cx);
447 tool_registry.register_tool(tools::now_tool::NowTool);
448}
449
450pub fn humanize_token_count(count: usize) -> String {
451 match count {
452 0..=999 => count.to_string(),
453 1000..=9999 => {
454 let thousands = count / 1000;
455 let hundreds = (count % 1000 + 50) / 100;
456 if hundreds == 0 {
457 format!("{}k", thousands)
458 } else if hundreds == 10 {
459 format!("{}k", thousands + 1)
460 } else {
461 format!("{}.{}k", thousands, hundreds)
462 }
463 }
464 _ => format!("{}k", (count + 500) / 1000),
465 }
466}
467
468#[cfg(test)]
469#[ctor::ctor]
470fn init_logger() {
471 if std::env::var("RUST_LOG").is_ok() {
472 env_logger::init();
473 }
474}