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 CycleMessageRole,
62 QuoteSelection,
63 InsertIntoEditor,
64 ToggleFocus,
65 InsertActivePrompt,
66 DeployHistory,
67 DeployPromptLibrary,
68 ConfirmCommand,
69 NewContext,
70 ToggleModelSelector,
71 ]
72);
73
74#[derive(PartialEq, Clone, Deserialize)]
75pub enum InsertDraggedFiles {
76 ProjectPaths(Vec<PathBuf>),
77 ExternalFiles(Vec<PathBuf>),
78}
79
80impl_actions!(assistant, [InsertDraggedFiles]);
81
82const DEFAULT_CONTEXT_LINES: usize = 50;
83
84#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
85pub struct MessageId(clock::Lamport);
86
87impl MessageId {
88 pub fn as_u64(self) -> u64 {
89 self.0.as_u64()
90 }
91}
92
93#[derive(Deserialize, Debug)]
94pub struct LanguageModelUsage {
95 pub prompt_tokens: u32,
96 pub completion_tokens: u32,
97 pub total_tokens: u32,
98}
99
100#[derive(Deserialize, Debug)]
101pub struct LanguageModelChoiceDelta {
102 pub index: u32,
103 pub delta: LanguageModelResponseMessage,
104 pub finish_reason: Option<String>,
105}
106
107#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
108pub enum MessageStatus {
109 Pending,
110 Done,
111 Error(SharedString),
112 Canceled,
113}
114
115impl MessageStatus {
116 pub fn from_proto(status: proto::ContextMessageStatus) -> MessageStatus {
117 match status.variant {
118 Some(proto::context_message_status::Variant::Pending(_)) => MessageStatus::Pending,
119 Some(proto::context_message_status::Variant::Done(_)) => MessageStatus::Done,
120 Some(proto::context_message_status::Variant::Error(error)) => {
121 MessageStatus::Error(error.message.into())
122 }
123 Some(proto::context_message_status::Variant::Canceled(_)) => MessageStatus::Canceled,
124 None => MessageStatus::Pending,
125 }
126 }
127
128 pub fn to_proto(&self) -> proto::ContextMessageStatus {
129 match self {
130 MessageStatus::Pending => proto::ContextMessageStatus {
131 variant: Some(proto::context_message_status::Variant::Pending(
132 proto::context_message_status::Pending {},
133 )),
134 },
135 MessageStatus::Done => proto::ContextMessageStatus {
136 variant: Some(proto::context_message_status::Variant::Done(
137 proto::context_message_status::Done {},
138 )),
139 },
140 MessageStatus::Error(message) => proto::ContextMessageStatus {
141 variant: Some(proto::context_message_status::Variant::Error(
142 proto::context_message_status::Error {
143 message: message.to_string(),
144 },
145 )),
146 },
147 MessageStatus::Canceled => proto::ContextMessageStatus {
148 variant: Some(proto::context_message_status::Variant::Canceled(
149 proto::context_message_status::Canceled {},
150 )),
151 },
152 }
153 }
154}
155
156/// The state pertaining to the Assistant.
157#[derive(Default)]
158struct Assistant {
159 /// Whether the Assistant is enabled.
160 enabled: bool,
161}
162
163impl Global for Assistant {}
164
165impl Assistant {
166 const NAMESPACE: &'static str = "assistant";
167
168 fn set_enabled(&mut self, enabled: bool, cx: &mut AppContext) {
169 if self.enabled == enabled {
170 return;
171 }
172
173 self.enabled = enabled;
174
175 if !enabled {
176 CommandPaletteFilter::update_global(cx, |filter, _cx| {
177 filter.hide_namespace(Self::NAMESPACE);
178 });
179
180 return;
181 }
182
183 CommandPaletteFilter::update_global(cx, |filter, _cx| {
184 filter.show_namespace(Self::NAMESPACE);
185 });
186 }
187}
188
189pub fn init(
190 fs: Arc<dyn Fs>,
191 client: Arc<Client>,
192 stdout_is_a_pty: bool,
193 cx: &mut AppContext,
194) -> Arc<PromptBuilder> {
195 cx.set_global(Assistant::default());
196 AssistantSettings::register(cx);
197 SlashCommandSettings::register(cx);
198
199 // TODO: remove this when 0.148.0 is released.
200 if AssistantSettings::get_global(cx).using_outdated_settings_version {
201 update_settings_file::<AssistantSettings>(fs.clone(), cx, {
202 let fs = fs.clone();
203 |content, cx| {
204 content.update_file(fs, cx);
205 }
206 });
207 }
208
209 cx.spawn(|mut cx| {
210 let client = client.clone();
211 async move {
212 let embedding_provider = CloudEmbeddingProvider::new(client.clone());
213 let semantic_index = SemanticDb::new(
214 paths::embeddings_dir().join("semantic-index-db.0.mdb"),
215 Arc::new(embedding_provider),
216 &mut cx,
217 )
218 .await?;
219
220 cx.update(|cx| cx.set_global(semantic_index))
221 }
222 })
223 .detach();
224
225 context_store::init(&client.clone().into());
226 prompt_library::init(cx);
227 init_language_model_settings(cx);
228 assistant_slash_command::init(cx);
229 assistant_tool::init(cx);
230 assistant_panel::init(cx);
231 context_servers::init(cx);
232
233 let prompt_builder = prompts::PromptBuilder::new(Some(PromptLoadingParams {
234 fs: fs.clone(),
235 repo_path: stdout_is_a_pty
236 .then(|| std::env::current_dir().log_err())
237 .flatten(),
238 cx,
239 }))
240 .log_err()
241 .map(Arc::new)
242 .unwrap_or_else(|| Arc::new(prompts::PromptBuilder::new(None).unwrap()));
243 register_slash_commands(Some(prompt_builder.clone()), cx);
244 register_tools(cx);
245 inline_assistant::init(
246 fs.clone(),
247 prompt_builder.clone(),
248 client.telemetry().clone(),
249 cx,
250 );
251 terminal_inline_assistant::init(
252 fs.clone(),
253 prompt_builder.clone(),
254 client.telemetry().clone(),
255 cx,
256 );
257 IndexedDocsRegistry::init_global(cx);
258
259 CommandPaletteFilter::update_global(cx, |filter, _cx| {
260 filter.hide_namespace(Assistant::NAMESPACE);
261 });
262 Assistant::update_global(cx, |assistant, cx| {
263 let settings = AssistantSettings::get_global(cx);
264
265 assistant.set_enabled(settings.enabled, cx);
266 });
267 cx.observe_global::<SettingsStore>(|cx| {
268 Assistant::update_global(cx, |assistant, cx| {
269 let settings = AssistantSettings::get_global(cx);
270 assistant.set_enabled(settings.enabled, cx);
271 });
272 })
273 .detach();
274
275 register_context_server_handlers(cx);
276
277 prompt_builder
278}
279
280fn register_context_server_handlers(cx: &mut AppContext) {
281 cx.subscribe(
282 &context_servers::manager::ContextServerManager::global(cx),
283 |manager, event, cx| match event {
284 context_servers::manager::Event::ServerStarted { server_id } => {
285 cx.update_model(
286 &manager,
287 |manager: &mut context_servers::manager::ContextServerManager, cx| {
288 let slash_command_registry = SlashCommandRegistry::global(cx);
289 let context_server_registry = ContextServerRegistry::global(cx);
290 if let Some(server) = manager.get_server(server_id) {
291 cx.spawn(|_, _| async move {
292 let Some(protocol) = server.client.read().clone() else {
293 return;
294 };
295
296 if let Some(prompts) = protocol.list_prompts().await.log_err() {
297 for prompt in prompts
298 .into_iter()
299 .filter(context_server_command::acceptable_prompt)
300 {
301 log::info!(
302 "registering context server command: {:?}",
303 prompt.name
304 );
305 context_server_registry.register_command(
306 server.id.clone(),
307 prompt.name.as_str(),
308 );
309 slash_command_registry.register_command(
310 context_server_command::ContextServerSlashCommand::new(
311 &server, prompt,
312 ),
313 true,
314 );
315 }
316 }
317 })
318 .detach();
319 }
320 },
321 );
322 }
323 context_servers::manager::Event::ServerStopped { server_id } => {
324 let slash_command_registry = SlashCommandRegistry::global(cx);
325 let context_server_registry = ContextServerRegistry::global(cx);
326 if let Some(commands) = context_server_registry.get_commands(server_id) {
327 for command_name in commands {
328 slash_command_registry.unregister_command_by_name(&command_name);
329 context_server_registry.unregister_command(&server_id, &command_name);
330 }
331 }
332 }
333 },
334 )
335 .detach();
336}
337
338fn init_language_model_settings(cx: &mut AppContext) {
339 update_active_language_model_from_settings(cx);
340
341 cx.observe_global::<SettingsStore>(update_active_language_model_from_settings)
342 .detach();
343 cx.subscribe(
344 &LanguageModelRegistry::global(cx),
345 |_, event: &language_model::Event, cx| match event {
346 language_model::Event::ProviderStateChanged
347 | language_model::Event::AddedProvider(_)
348 | language_model::Event::RemovedProvider(_) => {
349 update_active_language_model_from_settings(cx);
350 }
351 _ => {}
352 },
353 )
354 .detach();
355}
356
357fn update_active_language_model_from_settings(cx: &mut AppContext) {
358 let settings = AssistantSettings::get_global(cx);
359 let provider_name = LanguageModelProviderId::from(settings.default_model.provider.clone());
360 let model_id = LanguageModelId::from(settings.default_model.model.clone());
361 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
362 registry.select_active_model(&provider_name, &model_id, cx);
363 });
364}
365
366fn register_slash_commands(prompt_builder: Option<Arc<PromptBuilder>>, cx: &mut AppContext) {
367 let slash_command_registry = SlashCommandRegistry::global(cx);
368
369 slash_command_registry.register_command(file_command::FileSlashCommand, true);
370 slash_command_registry.register_command(delta_command::DeltaSlashCommand, true);
371 slash_command_registry.register_command(symbols_command::OutlineSlashCommand, true);
372 slash_command_registry.register_command(tab_command::TabSlashCommand, true);
373 slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
374 slash_command_registry.register_command(prompt_command::PromptSlashCommand, true);
375 slash_command_registry.register_command(default_command::DefaultSlashCommand, false);
376 slash_command_registry.register_command(terminal_command::TerminalSlashCommand, true);
377 slash_command_registry.register_command(now_command::NowSlashCommand, false);
378 slash_command_registry.register_command(diagnostics_command::DiagnosticsSlashCommand, true);
379
380 if let Some(prompt_builder) = prompt_builder {
381 slash_command_registry.register_command(
382 workflow_command::WorkflowSlashCommand::new(prompt_builder.clone()),
383 true,
384 );
385 }
386 slash_command_registry.register_command(fetch_command::FetchSlashCommand, false);
387
388 cx.observe_flag::<auto_command::AutoSlashCommandFeatureFlag, _>({
389 let slash_command_registry = slash_command_registry.clone();
390 move |is_enabled, _cx| {
391 if is_enabled {
392 // [#auto-staff-ship] TODO remove this when /auto is no longer staff-shipped
393 slash_command_registry.register_command(auto_command::AutoCommand, true);
394 }
395 }
396 })
397 .detach();
398
399 update_slash_commands_from_settings(cx);
400 cx.observe_global::<SettingsStore>(update_slash_commands_from_settings)
401 .detach();
402
403 cx.observe_flag::<search_command::SearchSlashCommandFeatureFlag, _>({
404 let slash_command_registry = slash_command_registry.clone();
405 move |is_enabled, _cx| {
406 if is_enabled {
407 slash_command_registry.register_command(search_command::SearchSlashCommand, true);
408 }
409 }
410 })
411 .detach();
412}
413
414fn update_slash_commands_from_settings(cx: &mut AppContext) {
415 let slash_command_registry = SlashCommandRegistry::global(cx);
416 let settings = SlashCommandSettings::get_global(cx);
417
418 if settings.docs.enabled {
419 slash_command_registry.register_command(docs_command::DocsSlashCommand, true);
420 } else {
421 slash_command_registry.unregister_command(docs_command::DocsSlashCommand);
422 }
423
424 if settings.project.enabled {
425 slash_command_registry.register_command(project_command::ProjectSlashCommand, true);
426 } else {
427 slash_command_registry.unregister_command(project_command::ProjectSlashCommand);
428 }
429}
430
431fn register_tools(cx: &mut AppContext) {
432 let tool_registry = ToolRegistry::global(cx);
433 tool_registry.register_tool(tools::now_tool::NowTool);
434}
435
436pub fn humanize_token_count(count: usize) -> String {
437 match count {
438 0..=999 => count.to_string(),
439 1000..=9999 => {
440 let thousands = count / 1000;
441 let hundreds = (count % 1000 + 50) / 100;
442 if hundreds == 0 {
443 format!("{}k", thousands)
444 } else if hundreds == 10 {
445 format!("{}k", thousands + 1)
446 } else {
447 format!("{}.{}k", thousands, hundreds)
448 }
449 }
450 _ => format!("{}k", (count + 500) / 1000),
451 }
452}
453
454#[cfg(test)]
455#[ctor::ctor]
456fn init_logger() {
457 if std::env::var("RUST_LOG").is_ok() {
458 env_logger::init();
459 }
460}