1mod assistant_settings;
2mod completion_provider;
3mod tools;
4mod 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
225struct 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 cx.view().downgrade(),
578 self.model.clone(),
579 body.clone(),
580 self.user_store.read(cx).current_user(),
581 self.can_submit(),
582 self.tool_registry.clone(),
583 ))
584 } else {
585 element
586 .on_click(cx.listener({
587 let id = *id;
588 move |assistant_chat, event: &ClickEvent, _cx| {
589 if event.up.click_count == 2 {
590 assistant_chat.editing_message_id = Some(id);
591 }
592 }
593 }))
594 .child(crate::ui::ChatMessage::new(
595 *id,
596 UserOrAssistant::User(self.user_store.read(cx).current_user()),
597 Some(
598 RichText::new(
599 body.read(cx).text(cx),
600 &[],
601 &self.language_registry,
602 )
603 .element(ElementId::from(id.0), cx),
604 ),
605 self.is_message_collapsed(id),
606 Box::new(cx.listener({
607 let id = *id;
608 move |assistant_chat, _event, _cx| {
609 assistant_chat.toggle_message_collapsed(id)
610 }
611 })),
612 ))
613 }
614 })
615 .into_any(),
616 ChatMessage::Assistant(AssistantMessage {
617 id,
618 body,
619 error,
620 tool_calls,
621 ..
622 }) => {
623 let assistant_body = if body.text.is_empty() {
624 None
625 } else {
626 Some(
627 div()
628 .p_2()
629 .child(body.element(ElementId::from(id.0), cx))
630 .into_any_element(),
631 )
632 };
633
634 div()
635 .when(!is_last, |element| element.mb_2())
636 .child(crate::ui::ChatMessage::new(
637 *id,
638 UserOrAssistant::Assistant,
639 assistant_body,
640 self.is_message_collapsed(id),
641 Box::new(cx.listener({
642 let id = *id;
643 move |assistant_chat, _event, _cx| {
644 assistant_chat.toggle_message_collapsed(id)
645 }
646 })),
647 ))
648 // TODO: Should the errors and tool calls get passed into `ChatMessage`?
649 .child(self.render_error(error.clone(), ix, cx))
650 .children(tool_calls.iter().map(|tool_call| {
651 let result = &tool_call.result;
652 let name = tool_call.name.clone();
653 match result {
654 Some(result) => {
655 div().p_2().child(result.into_any_element(&name)).into_any()
656 }
657 None => div()
658 .p_2()
659 .child(Label::new(name).color(Color::Modified))
660 .child("Running...")
661 .into_any(),
662 }
663 }))
664 .into_any()
665 }
666 }
667 }
668
669 fn completion_messages(&self, cx: &mut WindowContext) -> Vec<CompletionMessage> {
670 let mut completion_messages = Vec::new();
671
672 for message in &self.messages {
673 match message {
674 ChatMessage::User(UserMessage { body, .. }) => {
675 // When we re-introduce contexts like active file, we'll inject them here instead of relying on the model to request them
676 // contexts.iter().for_each(|context| {
677 // completion_messages.extend(context.completion_messages(cx))
678 // });
679
680 // Show user's message last so that the assistant is grounded in the user's request
681 completion_messages.push(CompletionMessage::User {
682 content: body.read(cx).text(cx),
683 });
684 }
685 ChatMessage::Assistant(AssistantMessage {
686 body, tool_calls, ..
687 }) => {
688 // In no case do we want to send an empty message. This shouldn't happen, but we might as well
689 // not break the Chat API if it does.
690 if body.text.is_empty() && tool_calls.is_empty() {
691 continue;
692 }
693
694 let tool_calls_from_assistant = tool_calls
695 .iter()
696 .map(|tool_call| ToolCall {
697 content: ToolCallContent::Function {
698 function: FunctionContent {
699 name: tool_call.name.clone(),
700 arguments: tool_call.arguments.clone(),
701 },
702 },
703 id: tool_call.id.clone(),
704 })
705 .collect();
706
707 completion_messages.push(CompletionMessage::Assistant {
708 content: Some(body.text.to_string()),
709 tool_calls: tool_calls_from_assistant,
710 });
711
712 for tool_call in tool_calls {
713 // todo!(): we should not be sending when the tool is still running / has no result
714 // For now I'm going to have to assume we send an empty string because otherwise
715 // the Chat API will break -- there is a required message for every tool call by ID
716 let content = match &tool_call.result {
717 Some(result) => result.format(&tool_call.name),
718 None => "".to_string(),
719 };
720
721 completion_messages.push(CompletionMessage::Tool {
722 content,
723 tool_call_id: tool_call.id.clone(),
724 });
725 }
726 }
727 }
728 }
729
730 completion_messages
731 }
732}
733
734impl Render for AssistantChat {
735 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
736 div()
737 .relative()
738 .flex_1()
739 .v_flex()
740 .key_context("AssistantChat")
741 .on_action(cx.listener(Self::submit))
742 .on_action(cx.listener(Self::cancel))
743 .on_action(cx.listener(Self::debug_project_index))
744 .text_color(Color::Default.color(cx))
745 .child(list(self.list_state.clone()).flex_1())
746 .child(Composer::new(
747 cx.view().downgrade(),
748 self.model.clone(),
749 self.composer_editor.clone(),
750 self.user_store.read(cx).current_user(),
751 self.can_submit(),
752 self.tool_registry.clone(),
753 ))
754 }
755}
756
757#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
758struct MessageId(usize);
759
760impl MessageId {
761 fn post_inc(&mut self) -> Self {
762 let id = *self;
763 self.0 += 1;
764 id
765 }
766}
767
768enum ChatMessage {
769 User(UserMessage),
770 Assistant(AssistantMessage),
771}
772
773impl ChatMessage {
774 fn focus_handle(&self, cx: &AppContext) -> Option<FocusHandle> {
775 match self {
776 ChatMessage::User(UserMessage { body, .. }) => Some(body.focus_handle(cx)),
777 ChatMessage::Assistant(_) => None,
778 }
779 }
780}
781
782struct UserMessage {
783 id: MessageId,
784 body: View<Editor>,
785}
786
787struct AssistantMessage {
788 id: MessageId,
789 body: RichText,
790 tool_calls: Vec<ToolFunctionCall>,
791 error: Option<SharedString>,
792}