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