1mod assistant_settings;
2mod completion_provider;
3mod tools;
4pub mod ui;
5
6use ::ui::{div, prelude::*, Color, ViewContext};
7use anyhow::{Context, Result};
8use assistant_tooling::{ToolFunctionCall, ToolRegistry};
9use client::{proto, Client, UserStore};
10use collections::HashMap;
11use completion_provider::*;
12use editor::Editor;
13use feature_flags::FeatureFlagAppExt as _;
14use futures::{future::join_all, StreamExt};
15use gpui::{
16 list, AnyElement, AppContext, AsyncWindowContext, ClickEvent, EventEmitter, FocusHandle,
17 FocusableView, ListAlignment, ListState, Model, Render, Task, View, WeakView,
18};
19use language::{language_settings::SoftWrap, LanguageRegistry};
20use open_ai::{FunctionContent, ToolCall, ToolCallContent};
21use rich_text::RichText;
22use semantic_index::{CloudEmbeddingProvider, ProjectIndex, SemanticIndex};
23use serde::Deserialize;
24use settings::Settings;
25use std::sync::Arc;
26use ui::Composer;
27use util::{paths::EMBEDDINGS_DIR, ResultExt};
28use workspace::{
29 dock::{DockPosition, Panel, PanelEvent},
30 Workspace,
31};
32
33pub use assistant_settings::AssistantSettings;
34
35use crate::tools::{CreateBufferTool, ProjectIndexTool};
36use crate::ui::UserOrAssistant;
37
38const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
39
40#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
41pub struct Submit(SubmitMode);
42
43/// There are multiple different ways to submit a model request, represented by this enum.
44#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
45pub enum SubmitMode {
46 /// Only include the conversation.
47 Simple,
48 /// Send the current file as context.
49 CurrentFile,
50 /// Search the codebase and send relevant excerpts.
51 Codebase,
52}
53
54gpui::actions!(assistant2, [Cancel, ToggleFocus, DebugProjectIndex]);
55gpui::impl_actions!(assistant2, [Submit]);
56
57pub fn init(client: Arc<Client>, cx: &mut AppContext) {
58 AssistantSettings::register(cx);
59
60 cx.spawn(|mut cx| {
61 let client = client.clone();
62 async move {
63 let embedding_provider = CloudEmbeddingProvider::new(client.clone());
64 let semantic_index = SemanticIndex::new(
65 EMBEDDINGS_DIR.join("semantic-index-db.0.mdb"),
66 Arc::new(embedding_provider),
67 &mut cx,
68 )
69 .await?;
70 cx.update(|cx| cx.set_global(semantic_index))
71 }
72 })
73 .detach();
74
75 cx.set_global(CompletionProvider::new(CloudCompletionProvider::new(
76 client,
77 )));
78
79 cx.observe_new_views(
80 |workspace: &mut Workspace, _cx: &mut ViewContext<Workspace>| {
81 workspace.register_action(|workspace, _: &ToggleFocus, cx| {
82 workspace.toggle_panel_focus::<AssistantPanel>(cx);
83 });
84 },
85 )
86 .detach();
87}
88
89pub fn enabled(cx: &AppContext) -> bool {
90 cx.is_staff()
91}
92
93pub struct AssistantPanel {
94 chat: View<AssistantChat>,
95 width: Option<Pixels>,
96}
97
98impl AssistantPanel {
99 pub fn load(
100 workspace: WeakView<Workspace>,
101 cx: AsyncWindowContext,
102 ) -> Task<Result<View<Self>>> {
103 cx.spawn(|mut cx| async move {
104 let (app_state, project) = workspace.update(&mut cx, |workspace, _| {
105 (workspace.app_state().clone(), workspace.project().clone())
106 })?;
107
108 let user_store = app_state.user_store.clone();
109
110 cx.new_view(|cx| {
111 // todo!("this will panic if the semantic index failed to load or has not loaded yet")
112 let project_index = cx.update_global(|semantic_index: &mut SemanticIndex, cx| {
113 semantic_index.project_index(project.clone(), cx)
114 });
115
116 let mut tool_registry = ToolRegistry::new();
117 tool_registry
118 .register(
119 ProjectIndexTool::new(project_index.clone(), app_state.fs.clone()),
120 cx,
121 )
122 .context("failed to register ProjectIndexTool")
123 .log_err();
124 tool_registry
125 .register(
126 CreateBufferTool::new(workspace.clone(), project.clone()),
127 cx,
128 )
129 .context("failed to register CreateBufferTool")
130 .log_err();
131
132 let tool_registry = Arc::new(tool_registry);
133
134 Self::new(
135 app_state.languages.clone(),
136 tool_registry,
137 user_store,
138 Some(project_index),
139 cx,
140 )
141 })
142 })
143 }
144
145 pub fn new(
146 language_registry: Arc<LanguageRegistry>,
147 tool_registry: Arc<ToolRegistry>,
148 user_store: Model<UserStore>,
149 project_index: Option<Model<ProjectIndex>>,
150 cx: &mut ViewContext<Self>,
151 ) -> Self {
152 let chat = cx.new_view(|cx| {
153 AssistantChat::new(
154 language_registry.clone(),
155 tool_registry.clone(),
156 user_store,
157 project_index,
158 cx,
159 )
160 });
161
162 Self { width: None, chat }
163 }
164}
165
166impl Render for AssistantPanel {
167 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
168 div()
169 .size_full()
170 .v_flex()
171 .p_2()
172 .bg(cx.theme().colors().background)
173 .child(self.chat.clone())
174 }
175}
176
177impl Panel for AssistantPanel {
178 fn persistent_name() -> &'static str {
179 "AssistantPanelv2"
180 }
181
182 fn position(&self, _cx: &WindowContext) -> workspace::dock::DockPosition {
183 // todo!("Add a setting / use assistant settings")
184 DockPosition::Right
185 }
186
187 fn position_is_valid(&self, position: workspace::dock::DockPosition) -> bool {
188 matches!(position, DockPosition::Right)
189 }
190
191 fn set_position(&mut self, _: workspace::dock::DockPosition, _: &mut ViewContext<Self>) {
192 // Do nothing until we have a setting for this
193 }
194
195 fn size(&self, _cx: &WindowContext) -> Pixels {
196 self.width.unwrap_or(px(400.))
197 }
198
199 fn set_size(&mut self, size: Option<Pixels>, cx: &mut ViewContext<Self>) {
200 self.width = size;
201 cx.notify();
202 }
203
204 fn icon(&self, _cx: &WindowContext) -> Option<::ui::IconName> {
205 Some(IconName::Ai)
206 }
207
208 fn icon_tooltip(&self, _: &WindowContext) -> Option<&'static str> {
209 Some("Assistant Panel ✨")
210 }
211
212 fn toggle_action(&self) -> Box<dyn gpui::Action> {
213 Box::new(ToggleFocus)
214 }
215}
216
217impl EventEmitter<PanelEvent> for AssistantPanel {}
218
219impl FocusableView for AssistantPanel {
220 fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
221 self.chat.read(cx).composer_editor.read(cx).focus_handle(cx)
222 }
223}
224
225pub struct AssistantChat {
226 model: String,
227 messages: Vec<ChatMessage>,
228 list_state: ListState,
229 language_registry: Arc<LanguageRegistry>,
230 composer_editor: View<Editor>,
231 user_store: Model<UserStore>,
232 next_message_id: MessageId,
233 collapsed_messages: HashMap<MessageId, bool>,
234 editing_message_id: Option<MessageId>,
235 pending_completion: Option<Task<()>>,
236 tool_registry: Arc<ToolRegistry>,
237 project_index: Option<Model<ProjectIndex>>,
238}
239
240impl AssistantChat {
241 fn new(
242 language_registry: Arc<LanguageRegistry>,
243 tool_registry: Arc<ToolRegistry>,
244 user_store: Model<UserStore>,
245 project_index: Option<Model<ProjectIndex>>,
246 cx: &mut ViewContext<Self>,
247 ) -> Self {
248 let model = CompletionProvider::get(cx).default_model();
249 let view = cx.view().downgrade();
250 let list_state = ListState::new(
251 0,
252 ListAlignment::Bottom,
253 px(1024.),
254 move |ix, cx: &mut WindowContext| {
255 view.update(cx, |this, cx| this.render_message(ix, cx))
256 .unwrap()
257 },
258 );
259
260 Self {
261 model,
262 messages: Vec::new(),
263 composer_editor: cx.new_view(|cx| {
264 let mut editor = Editor::auto_height(80, cx);
265 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
266 editor.set_placeholder_text("Type a message to the assistant", cx);
267 editor
268 }),
269 list_state,
270 user_store,
271 language_registry,
272 project_index,
273 next_message_id: MessageId(0),
274 editing_message_id: None,
275 collapsed_messages: HashMap::default(),
276 pending_completion: None,
277 tool_registry,
278 }
279 }
280
281 fn focused_message_id(&self, cx: &WindowContext) -> Option<MessageId> {
282 self.messages.iter().find_map(|message| match message {
283 ChatMessage::User(message) => message
284 .body
285 .focus_handle(cx)
286 .contains_focused(cx)
287 .then_some(message.id),
288 ChatMessage::Assistant(_) => None,
289 })
290 }
291
292 fn cancel(&mut self, _: &Cancel, cx: &mut ViewContext<Self>) {
293 // If we're currently editing a message, cancel the edit.
294 self.editing_message_id.take();
295
296 if self.pending_completion.take().is_none() {
297 cx.propagate();
298 return;
299 }
300
301 if let Some(ChatMessage::Assistant(message)) = self.messages.last() {
302 if message.body.text.is_empty() {
303 self.pop_message(cx);
304 }
305 }
306 }
307
308 fn submit(&mut self, Submit(mode): &Submit, cx: &mut ViewContext<Self>) {
309 // Don't allow multiple concurrent completions.
310 if self.pending_completion.is_some() {
311 cx.propagate();
312 return;
313 }
314
315 if let Some(focused_message_id) = self.focused_message_id(cx) {
316 self.truncate_messages(focused_message_id, cx);
317 } else if self.composer_editor.focus_handle(cx).is_focused(cx) {
318 let message = self.composer_editor.update(cx, |composer_editor, cx| {
319 let text = composer_editor.text(cx);
320 let id = self.next_message_id.post_inc();
321 let body = cx.new_view(|cx| {
322 let mut editor = Editor::auto_height(80, cx);
323 editor.set_text(text, cx);
324 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
325 editor
326 });
327 composer_editor.clear(cx);
328 ChatMessage::User(UserMessage { id, body })
329 });
330 self.push_message(message, cx);
331 } else {
332 log::error!("unexpected state: no user message editor is focused.");
333 return;
334 }
335
336 let mode = *mode;
337 self.pending_completion = Some(cx.spawn(move |this, mut cx| async move {
338 Self::request_completion(
339 this.clone(),
340 mode,
341 MAX_COMPLETION_CALLS_PER_SUBMISSION,
342 &mut cx,
343 )
344 .await
345 .log_err();
346
347 this.update(&mut cx, |this, cx| {
348 let composer_focus_handle = this.composer_editor.focus_handle(cx);
349 cx.focus(&composer_focus_handle);
350 this.pending_completion = None;
351 })
352 .context("Failed to push new user message")
353 .log_err();
354 }));
355 }
356
357 fn can_submit(&self) -> bool {
358 self.pending_completion.is_none()
359 }
360
361 fn debug_project_index(&mut self, _: &DebugProjectIndex, cx: &mut ViewContext<Self>) {
362 if let Some(index) = &self.project_index {
363 index.update(cx, |project_index, cx| {
364 project_index.debug(cx).detach_and_log_err(cx)
365 });
366 }
367 }
368
369 async fn request_completion(
370 this: WeakView<Self>,
371 mode: SubmitMode,
372 limit: usize,
373 cx: &mut AsyncWindowContext,
374 ) -> Result<()> {
375 let mut call_count = 0;
376 loop {
377 let complete = async {
378 let completion = this.update(cx, |this, cx| {
379 this.push_new_assistant_message(cx);
380
381 let definitions = if call_count < limit
382 && matches!(mode, SubmitMode::Codebase | SubmitMode::Simple)
383 {
384 this.tool_registry.definitions()
385 } else {
386 &[]
387 };
388 call_count += 1;
389
390 let messages = this.completion_messages(cx);
391
392 CompletionProvider::get(cx).complete(
393 this.model.clone(),
394 messages,
395 Vec::new(),
396 1.0,
397 definitions,
398 )
399 });
400
401 let mut stream = completion?.await?;
402 let mut body = String::new();
403 while let Some(delta) = stream.next().await {
404 let delta = delta?;
405 this.update(cx, |this, cx| {
406 if let Some(ChatMessage::Assistant(AssistantMessage {
407 body: message_body,
408 tool_calls: message_tool_calls,
409 ..
410 })) = this.messages.last_mut()
411 {
412 if let Some(content) = &delta.content {
413 body.push_str(content);
414 }
415
416 for tool_call in delta.tool_calls {
417 let index = tool_call.index as usize;
418 if index >= message_tool_calls.len() {
419 message_tool_calls.resize_with(index + 1, Default::default);
420 }
421 let call = &mut message_tool_calls[index];
422
423 if let Some(id) = &tool_call.id {
424 call.id.push_str(id);
425 }
426
427 match tool_call.variant {
428 Some(proto::tool_call_delta::Variant::Function(tool_call)) => {
429 if let Some(name) = &tool_call.name {
430 call.name.push_str(name);
431 }
432 if let Some(arguments) = &tool_call.arguments {
433 call.arguments.push_str(arguments);
434 }
435 }
436 None => {}
437 }
438 }
439
440 *message_body =
441 RichText::new(body.clone(), &[], &this.language_registry);
442 cx.notify();
443 } else {
444 unreachable!()
445 }
446 })?;
447 }
448
449 anyhow::Ok(())
450 }
451 .await;
452
453 let mut tool_tasks = Vec::new();
454 this.update(cx, |this, cx| {
455 if let Some(ChatMessage::Assistant(AssistantMessage {
456 error: message_error,
457 tool_calls,
458 ..
459 })) = this.messages.last_mut()
460 {
461 if let Err(error) = complete {
462 message_error.replace(SharedString::from(error.to_string()));
463 cx.notify();
464 } else {
465 for tool_call in tool_calls.iter() {
466 tool_tasks.push(this.tool_registry.call(tool_call, cx));
467 }
468 }
469 }
470 })?;
471
472 if tool_tasks.is_empty() {
473 return Ok(());
474 }
475
476 let tools = join_all(tool_tasks.into_iter()).await;
477 // If the WindowContext went away for any tool's view we don't include it
478 // especially since the below call would fail for the same reason.
479 let tools = tools.into_iter().filter_map(|tool| tool.ok()).collect();
480
481 this.update(cx, |this, cx| {
482 if let Some(ChatMessage::Assistant(AssistantMessage { tool_calls, .. })) =
483 this.messages.last_mut()
484 {
485 *tool_calls = tools;
486 cx.notify();
487 }
488 })?;
489 }
490 }
491
492 fn push_new_assistant_message(&mut self, cx: &mut ViewContext<Self>) {
493 let message = ChatMessage::Assistant(AssistantMessage {
494 id: self.next_message_id.post_inc(),
495 body: RichText::default(),
496 tool_calls: Vec::new(),
497 error: None,
498 });
499 self.push_message(message, cx);
500 }
501
502 fn push_message(&mut self, message: ChatMessage, cx: &mut ViewContext<Self>) {
503 let old_len = self.messages.len();
504 let focus_handle = Some(message.focus_handle(cx));
505 self.messages.push(message);
506 self.list_state
507 .splice_focusable(old_len..old_len, focus_handle);
508 cx.notify();
509 }
510
511 fn pop_message(&mut self, cx: &mut ViewContext<Self>) {
512 if self.messages.is_empty() {
513 return;
514 }
515
516 self.messages.pop();
517 self.list_state
518 .splice(self.messages.len()..self.messages.len() + 1, 0);
519 cx.notify();
520 }
521
522 fn truncate_messages(&mut self, last_message_id: MessageId, cx: &mut ViewContext<Self>) {
523 if let Some(index) = self.messages.iter().position(|message| match message {
524 ChatMessage::User(message) => message.id == last_message_id,
525 ChatMessage::Assistant(message) => message.id == last_message_id,
526 }) {
527 self.list_state.splice(index + 1..self.messages.len(), 0);
528 self.messages.truncate(index + 1);
529 cx.notify();
530 }
531 }
532
533 fn is_message_collapsed(&self, id: &MessageId) -> bool {
534 self.collapsed_messages.get(id).copied().unwrap_or_default()
535 }
536
537 fn toggle_message_collapsed(&mut self, id: MessageId) {
538 let entry = self.collapsed_messages.entry(id).or_insert(false);
539 *entry = !*entry;
540 }
541
542 fn render_error(
543 &self,
544 error: Option<SharedString>,
545 _ix: usize,
546 cx: &mut ViewContext<Self>,
547 ) -> AnyElement {
548 let theme = cx.theme();
549
550 if let Some(error) = error {
551 div()
552 .py_1()
553 .px_2()
554 .neg_mx_1()
555 .rounded_md()
556 .border()
557 .border_color(theme.status().error_border)
558 // .bg(theme.status().error_background)
559 .text_color(theme.status().error)
560 .child(error.clone())
561 .into_any_element()
562 } else {
563 div().into_any_element()
564 }
565 }
566
567 fn render_message(&self, ix: usize, cx: &mut ViewContext<Self>) -> AnyElement {
568 let is_last = ix == self.messages.len() - 1;
569
570 match &self.messages[ix] {
571 ChatMessage::User(UserMessage { id, body }) => div()
572 .id(SharedString::from(format!("message-{}-container", id.0)))
573 .when(!is_last, |element| element.mb_2())
574 .map(|element| {
575 if self.editing_message_id.as_ref() == Some(id) {
576 element.child(Composer::new(
577 body.clone(),
578 self.user_store.read(cx).current_user(),
579 self.can_submit(),
580 self.tool_registry.clone(),
581 crate::ui::ModelSelector::new(
582 cx.view().downgrade(),
583 self.model.clone(),
584 )
585 .into_any_element(),
586 ))
587 } else {
588 element
589 .on_click(cx.listener({
590 let id = *id;
591 move |assistant_chat, event: &ClickEvent, _cx| {
592 if event.up.click_count == 2 {
593 assistant_chat.editing_message_id = Some(id);
594 }
595 }
596 }))
597 .child(crate::ui::ChatMessage::new(
598 *id,
599 UserOrAssistant::User(self.user_store.read(cx).current_user()),
600 Some(
601 RichText::new(
602 body.read(cx).text(cx),
603 &[],
604 &self.language_registry,
605 )
606 .element(ElementId::from(id.0), cx),
607 ),
608 self.is_message_collapsed(id),
609 Box::new(cx.listener({
610 let id = *id;
611 move |assistant_chat, _event, _cx| {
612 assistant_chat.toggle_message_collapsed(id)
613 }
614 })),
615 ))
616 }
617 })
618 .into_any(),
619 ChatMessage::Assistant(AssistantMessage {
620 id,
621 body,
622 error,
623 tool_calls,
624 ..
625 }) => {
626 let assistant_body = if body.text.is_empty() {
627 None
628 } else {
629 Some(
630 div()
631 .p_2()
632 .child(body.element(ElementId::from(id.0), cx))
633 .into_any_element(),
634 )
635 };
636
637 div()
638 .when(!is_last, |element| element.mb_2())
639 .child(crate::ui::ChatMessage::new(
640 *id,
641 UserOrAssistant::Assistant,
642 assistant_body,
643 self.is_message_collapsed(id),
644 Box::new(cx.listener({
645 let id = *id;
646 move |assistant_chat, _event, _cx| {
647 assistant_chat.toggle_message_collapsed(id)
648 }
649 })),
650 ))
651 // TODO: Should the errors and tool calls get passed into `ChatMessage`?
652 .child(self.render_error(error.clone(), ix, cx))
653 .children(tool_calls.iter().map(|tool_call| {
654 let result = &tool_call.result;
655 let name = tool_call.name.clone();
656 match result {
657 Some(result) => {
658 div().p_2().child(result.into_any_element(&name)).into_any()
659 }
660 None => div()
661 .p_2()
662 .child(Label::new(name).color(Color::Modified))
663 .child("Running...")
664 .into_any(),
665 }
666 }))
667 .into_any()
668 }
669 }
670 }
671
672 fn completion_messages(&self, cx: &mut WindowContext) -> Vec<CompletionMessage> {
673 let mut completion_messages = Vec::new();
674
675 for message in &self.messages {
676 match message {
677 ChatMessage::User(UserMessage { body, .. }) => {
678 // When we re-introduce contexts like active file, we'll inject them here instead of relying on the model to request them
679 // contexts.iter().for_each(|context| {
680 // completion_messages.extend(context.completion_messages(cx))
681 // });
682
683 // Show user's message last so that the assistant is grounded in the user's request
684 completion_messages.push(CompletionMessage::User {
685 content: body.read(cx).text(cx),
686 });
687 }
688 ChatMessage::Assistant(AssistantMessage {
689 body, tool_calls, ..
690 }) => {
691 // In no case do we want to send an empty message. This shouldn't happen, but we might as well
692 // not break the Chat API if it does.
693 if body.text.is_empty() && tool_calls.is_empty() {
694 continue;
695 }
696
697 let tool_calls_from_assistant = tool_calls
698 .iter()
699 .map(|tool_call| ToolCall {
700 content: ToolCallContent::Function {
701 function: FunctionContent {
702 name: tool_call.name.clone(),
703 arguments: tool_call.arguments.clone(),
704 },
705 },
706 id: tool_call.id.clone(),
707 })
708 .collect();
709
710 completion_messages.push(CompletionMessage::Assistant {
711 content: Some(body.text.to_string()),
712 tool_calls: tool_calls_from_assistant,
713 });
714
715 for tool_call in tool_calls {
716 // todo!(): we should not be sending when the tool is still running / has no result
717 // For now I'm going to have to assume we send an empty string because otherwise
718 // the Chat API will break -- there is a required message for every tool call by ID
719 let content = match &tool_call.result {
720 Some(result) => result.format(&tool_call.name),
721 None => "".to_string(),
722 };
723
724 completion_messages.push(CompletionMessage::Tool {
725 content,
726 tool_call_id: tool_call.id.clone(),
727 });
728 }
729 }
730 }
731 }
732
733 completion_messages
734 }
735}
736
737impl Render for AssistantChat {
738 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
739 div()
740 .relative()
741 .flex_1()
742 .v_flex()
743 .key_context("AssistantChat")
744 .on_action(cx.listener(Self::submit))
745 .on_action(cx.listener(Self::cancel))
746 .on_action(cx.listener(Self::debug_project_index))
747 .text_color(Color::Default.color(cx))
748 .child(list(self.list_state.clone()).flex_1())
749 .child(Composer::new(
750 self.composer_editor.clone(),
751 self.user_store.read(cx).current_user(),
752 self.can_submit(),
753 self.tool_registry.clone(),
754 crate::ui::ModelSelector::new(cx.view().downgrade(), self.model.clone())
755 .into_any_element(),
756 ))
757 }
758}
759
760#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
761pub struct MessageId(usize);
762
763impl MessageId {
764 fn post_inc(&mut self) -> Self {
765 let id = *self;
766 self.0 += 1;
767 id
768 }
769}
770
771enum ChatMessage {
772 User(UserMessage),
773 Assistant(AssistantMessage),
774}
775
776impl ChatMessage {
777 fn focus_handle(&self, cx: &AppContext) -> Option<FocusHandle> {
778 match self {
779 ChatMessage::User(UserMessage { body, .. }) => Some(body.focus_handle(cx)),
780 ChatMessage::Assistant(_) => None,
781 }
782 }
783}
784
785struct UserMessage {
786 id: MessageId,
787 body: View<Editor>,
788}
789
790struct AssistantMessage {
791 id: MessageId,
792 body: RichText,
793 tool_calls: Vec<ToolFunctionCall>,
794 error: Option<SharedString>,
795}