assistant2: Add rudimentary chat functionality (#21178)

This PR adds in rudimentary functionality for sending messages to the
LLM in `assistant2`.

<img width="1079" alt="Screenshot 2024-11-25 at 1 49 11 PM"
src="https://github.com/user-attachments/assets/5accb749-c034-4fb2-bf55-3ae5bc9529ad">

Release Notes:

- N/A

crates/assistant2/Cargo.toml

@@ -17,11 +17,14 @@ anyhow.workspace = true
 command_palette_hooks.workspace = true
 editor.workspace = true
 feature_flags.workspace = true
+futures.workspace = true
 gpui.workspace = true
 language_model.workspace = true
 language_model_selector.workspace = true
 proto.workspace = true
 settings.workspace = true
+smol.workspace = true
 theme.workspace = true
 ui.workspace = true
+util.workspace = true
 workspace.workspace = true

crates/assistant2/src/assistant.rs

@@ -1,5 +1,6 @@
 mod assistant_panel;
 mod message_editor;
+mod thread;
 
 use command_palette_hooks::CommandPaletteFilter;
 use feature_flags::{Assistant2FeatureFlag, FeatureFlagAppExt};
@@ -7,7 +8,10 @@ use gpui::{actions, AppContext};
 
 pub use crate::assistant_panel::AssistantPanel;
 
-actions!(assistant2, [ToggleFocus, NewThread, ToggleModelSelector]);
+actions!(
+    assistant2,
+    [ToggleFocus, NewThread, ToggleModelSelector, Chat]
+);
 
 const NAMESPACE: &str = "assistant2";
 

crates/assistant2/src/assistant_panel.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use gpui::{
     prelude::*, px, Action, AppContext, AsyncWindowContext, EventEmitter, FocusHandle,
-    FocusableView, Pixels, Task, View, ViewContext, WeakView, WindowContext,
+    FocusableView, Model, Pixels, Task, View, ViewContext, WeakView, WindowContext,
 };
 use language_model::LanguageModelRegistry;
 use language_model_selector::LanguageModelSelector;
@@ -10,6 +10,7 @@ use workspace::dock::{DockPosition, Panel, PanelEvent};
 use workspace::{Pane, Workspace};
 
 use crate::message_editor::MessageEditor;
+use crate::thread::Thread;
 use crate::{NewThread, ToggleFocus, ToggleModelSelector};
 
 pub fn init(cx: &mut AppContext) {
@@ -25,6 +26,7 @@ pub fn init(cx: &mut AppContext) {
 
 pub struct AssistantPanel {
     pane: View<Pane>,
+    thread: Model<Thread>,
     message_editor: View<MessageEditor>,
 }
 
@@ -56,9 +58,12 @@ impl AssistantPanel {
             pane
         });
 
+        let thread = cx.new_model(Thread::new);
+
         Self {
             pane,
-            message_editor: cx.new_view(MessageEditor::new),
+            thread: thread.clone(),
+            message_editor: cx.new_view(|cx| MessageEditor::new(thread, cx)),
         }
     }
 }
@@ -247,7 +252,24 @@ impl Render for AssistantPanel {
                 println!("Action: New Thread");
             }))
             .child(self.render_toolbar(cx))
-            .child(v_flex().bg(cx.theme().colors().panel_background))
+            .child(
+                v_flex()
+                    .id("message-list")
+                    .gap_2()
+                    .size_full()
+                    .p_2()
+                    .overflow_y_scroll()
+                    .bg(cx.theme().colors().panel_background)
+                    .children(self.thread.read(cx).messages.iter().map(|message| {
+                        v_flex()
+                            .p_2()
+                            .border_1()
+                            .border_color(cx.theme().colors().border_variant)
+                            .rounded_md()
+                            .child(Label::new(message.role.to_string()))
+                            .child(Label::new(message.text.clone()))
+                    })),
+            )
             .child(
                 h_flex()
                     .border_t_1()

crates/assistant2/src/message_editor.rs

@@ -1,16 +1,32 @@
 use editor::{Editor, EditorElement, EditorStyle};
-use gpui::{TextStyle, View};
+use futures::StreamExt;
+use gpui::{AppContext, Model, TextStyle, View};
+use language_model::{
+    LanguageModelCompletionEvent, LanguageModelRegistry, LanguageModelRequest,
+    LanguageModelRequestMessage, MessageContent, Role, StopReason,
+};
 use settings::Settings;
 use theme::ThemeSettings;
-use ui::prelude::*;
+use ui::{prelude::*, ButtonLike, ElevationIndex, KeyBinding};
+use util::ResultExt;
+
+use crate::thread::{self, Thread};
+use crate::Chat;
+
+#[derive(Debug, Clone, Copy)]
+pub enum RequestKind {
+    Chat,
+}
 
 pub struct MessageEditor {
+    thread: Model<Thread>,
     editor: View<Editor>,
 }
 
 impl MessageEditor {
-    pub fn new(cx: &mut ViewContext<Self>) -> Self {
+    pub fn new(thread: Model<Thread>, cx: &mut ViewContext<Self>) -> Self {
         Self {
+            thread,
             editor: cx.new_view(|cx| {
                 let mut editor = Editor::auto_height(80, cx);
                 editor.set_placeholder_text("Ask anything…", cx);
@@ -19,14 +35,122 @@ impl MessageEditor {
             }),
         }
     }
+
+    fn chat(&mut self, _: &Chat, cx: &mut ViewContext<Self>) {
+        self.send_to_model(RequestKind::Chat, cx);
+    }
+
+    fn send_to_model(
+        &mut self,
+        request_kind: RequestKind,
+        cx: &mut ViewContext<Self>,
+    ) -> Option<()> {
+        let provider = LanguageModelRegistry::read_global(cx).active_provider();
+        if provider
+            .as_ref()
+            .map_or(false, |provider| provider.must_accept_terms(cx))
+        {
+            cx.notify();
+            return None;
+        }
+
+        let model_registry = LanguageModelRegistry::read_global(cx);
+        let model = model_registry.active_model()?;
+
+        let request = self.build_completion_request(request_kind, cx);
+
+        let user_message = self.editor.read(cx).text(cx);
+        self.thread.update(cx, |thread, _cx| {
+            thread.messages.push(thread::Message {
+                role: Role::User,
+                text: user_message,
+            });
+        });
+
+        self.editor.update(cx, |editor, cx| {
+            editor.clear(cx);
+        });
+
+        let task = cx.spawn(|this, mut cx| async move {
+            let stream = model.stream_completion(request, &cx);
+            let stream_completion = async {
+                let mut events = stream.await?;
+                let mut stop_reason = StopReason::EndTurn;
+
+                let mut text = String::new();
+
+                while let Some(event) = events.next().await {
+                    let event = event?;
+                    match event {
+                        LanguageModelCompletionEvent::StartMessage { .. } => {}
+                        LanguageModelCompletionEvent::Stop(reason) => {
+                            stop_reason = reason;
+                        }
+                        LanguageModelCompletionEvent::Text(chunk) => {
+                            text.push_str(&chunk);
+                        }
+                        LanguageModelCompletionEvent::ToolUse(_tool_use) => {}
+                    }
+
+                    smol::future::yield_now().await;
+                }
+
+                anyhow::Ok((stop_reason, text))
+            };
+
+            let result = stream_completion.await;
+
+            this.update(&mut cx, |this, cx| {
+                if let Some((_stop_reason, text)) = result.log_err() {
+                    this.thread.update(cx, |thread, _cx| {
+                        thread.messages.push(thread::Message {
+                            role: Role::Assistant,
+                            text,
+                        });
+                    });
+                }
+            })
+            .ok();
+        });
+
+        self.thread.update(cx, |thread, _cx| {
+            thread.pending_completion_tasks.push(task);
+        });
+
+        None
+    }
+
+    fn build_completion_request(
+        &self,
+        _request_kind: RequestKind,
+        cx: &AppContext,
+    ) -> LanguageModelRequest {
+        let text = self.editor.read(cx).text(cx);
+
+        let request = LanguageModelRequest {
+            messages: vec![LanguageModelRequestMessage {
+                role: Role::User,
+                content: vec![MessageContent::Text(text)],
+                cache: false,
+            }],
+            tools: Vec::new(),
+            stop: Vec::new(),
+            temperature: None,
+        };
+
+        request
+    }
 }
 
 impl Render for MessageEditor {
     fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
         let font_size = TextSize::Default.rems(cx);
         let line_height = font_size.to_pixels(cx.rem_size()) * 1.3;
+        let focus_handle = self.editor.focus_handle(cx);
 
         v_flex()
+            .key_context("MessageEditor")
+            .on_action(cx.listener(Self::chat))
             .size_full()
             .gap_2()
             .p_2()
@@ -69,7 +193,19 @@ impl Render for MessageEditor {
                             .gap_2()
                             .child(Button::new("codebase", "Codebase").style(ButtonStyle::Filled))
                             .child(Label::new("or"))
-                            .child(Button::new("chat", "Chat").style(ButtonStyle::Filled)),
+                            .child(
+                                ButtonLike::new("chat")
+                                    .style(ButtonStyle::Filled)
+                                    .layer(ElevationIndex::ModalSurface)
+                                    .child(Label::new("Chat"))
+                                    .children(
+                                        KeyBinding::for_action_in(&Chat, &focus_handle, cx)
+                                            .map(|binding| binding.into_any_element()),
+                                    )
+                                    .on_click(move |_event, cx| {
+                                        focus_handle.dispatch_action(&Chat, cx);
+                                    }),
+                            ),
                     ),
             )
     }

crates/assistant2/src/thread.rs

@@ -0,0 +1,23 @@
+use gpui::{ModelContext, Task};
+use language_model::Role;
+
+/// A message in a [`Thread`].
+pub struct Message {
+    pub role: Role,
+    pub text: String,
+}
+
+/// A thread of conversation with the LLM.
+pub struct Thread {
+    pub messages: Vec<Message>,
+    pub pending_completion_tasks: Vec<Task<()>>,
+}
+
+impl Thread {
+    pub fn new(_cx: &mut ModelContext<Self>) -> Self {
+        Self {
+            messages: Vec::new(),
+            pending_completion_tasks: Vec::new(),
+        }
+    }
+}