From 7e418cc8afda7e54fa1098eb36664dbfe25863de Mon Sep 17 00:00:00 2001
From: Marshall Bowers <elliott.codes@gmail.com>
Date: Mon, 25 Nov 2024 20:49:03 -0500
Subject: [PATCH] assistant2: Style messages (#21191)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR styles the messages in `assistant2` so they don't look quite as
rough:

<img width="1138" alt="Screenshot 2024-11-25 at 8 36 32 PM"
src="https://github.com/user-attachments/assets/9cc423fa-feff-4c69-9d2b-d28970559547">

Release Notes:

- N/A
---
 crates/assistant2/src/assistant_panel.rs | 47 ++++++++++++++++++------
 crates/assistant2/src/thread.rs          |  2 +-
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/crates/assistant2/src/assistant_panel.rs b/crates/assistant2/src/assistant_panel.rs
index b05a39a1cd500657dfaeaa9dfc2bd2cc6ced586d..4ebf07e9d43bda37190cfc5acde48e7100848d8c 100644
--- a/crates/assistant2/src/assistant_panel.rs
+++ b/crates/assistant2/src/assistant_panel.rs
@@ -6,14 +6,14 @@ use gpui::{
     prelude::*, px, Action, AppContext, AsyncWindowContext, EventEmitter, FocusHandle,
     FocusableView, Model, Pixels, Subscription, Task, View, ViewContext, WeakView, WindowContext,
 };
-use language_model::LanguageModelRegistry;
+use language_model::{LanguageModelRegistry, Role};
 use language_model_selector::LanguageModelSelector;
 use ui::{prelude::*, ButtonLike, Divider, IconButtonShape, Tab, Tooltip};
 use workspace::dock::{DockPosition, Panel, PanelEvent};
 use workspace::Workspace;
 
 use crate::message_editor::MessageEditor;
-use crate::thread::{Thread, ThreadEvent};
+use crate::thread::{Message, Thread, ThreadEvent};
 use crate::{NewThread, ToggleFocus, ToggleModelSelector};
 
 pub fn init(cx: &mut AppContext) {
@@ -272,10 +272,39 @@ impl AssistantPanel {
                 .tooltip(move |cx| Tooltip::for_action("Change Model", &ToggleModelSelector, cx)),
         )
     }
+
+    fn render_message(&self, message: Message, cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let (role_icon, role_name) = match message.role {
+            Role::User => (IconName::Person, "You"),
+            Role::Assistant => (IconName::ZedAssistant, "Assistant"),
+            Role::System => (IconName::Settings, "System"),
+        };
+
+        v_flex()
+            .border_1()
+            .border_color(cx.theme().colors().border_variant)
+            .rounded_md()
+            .child(
+                h_flex()
+                    .justify_between()
+                    .p_1p5()
+                    .border_b_1()
+                    .border_color(cx.theme().colors().border_variant)
+                    .child(
+                        h_flex()
+                            .gap_2()
+                            .child(Icon::new(role_icon).size(IconSize::Small))
+                            .child(Label::new(role_name).size(LabelSize::Small)),
+                    ),
+            )
+            .child(v_flex().p_1p5().child(Label::new(message.text.clone())))
+    }
 }
 
 impl Render for AssistantPanel {
     fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let messages = self.thread.read(cx).messages().cloned().collect::<Vec<_>>();
+
         v_flex()
             .key_context("AssistantPanel2")
             .justify_between()
@@ -292,15 +321,11 @@ impl Render for AssistantPanel {
                     .p_2()
                     .overflow_y_scroll()
                     .bg(cx.theme().colors().panel_background)
-                    .children(self.thread.read(cx).messages().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()))
-                    })),
+                    .children(
+                        messages
+                            .into_iter()
+                            .map(|message| self.render_message(message, cx)),
+                    ),
             )
             .child(
                 h_flex()
diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs
index 067e82a60250ea8a2f6b22b792e291079e021307..d8263d15f7c45f860538ebda12ef84773d3657a9 100644
--- a/crates/assistant2/src/thread.rs
+++ b/crates/assistant2/src/thread.rs
@@ -18,7 +18,7 @@ pub enum RequestKind {
 }
 
 /// A message in a [`Thread`].
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Message {
     pub role: Role,
     pub text: String,