diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs index 5349a7ee86..0a5260baca 100644 --- a/crates/assistant2/src/assistant2.rs +++ b/crates/assistant2/src/assistant2.rs @@ -33,6 +33,8 @@ use workspace::{ pub use assistant_settings::AssistantSettings; +use crate::ui::{ChatMessageHeader, UserOrAssistant}; + const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5; #[derive(Eq, PartialEq, Copy, Clone, Deserialize)] @@ -526,7 +528,9 @@ impl AssistantChat { match &self.messages[ix] { ChatMessage::User(UserMessage { body, .. }) => div() .when(!is_last, |element| element.mb_2()) - .child(div().p_2().child(Label::new("You").color(Color::Default))) + .child(ChatMessageHeader::new(UserOrAssistant::User( + self.user_store.read(cx).current_user(), + ))) .child( div() .p_2() @@ -551,11 +555,7 @@ impl AssistantChat { div() .when(!is_last, |element| element.mb_2()) - .child( - div() - .p_2() - .child(Label::new("Assistant").color(Color::Modified)), - ) + .child(ChatMessageHeader::new(UserOrAssistant::Assistant)) .child(assistant_body) .child(self.render_error(error.clone(), ix, cx)) .children(tool_calls.iter().map(|tool_call| { diff --git a/crates/assistant2/src/ui.rs b/crates/assistant2/src/ui.rs index 5520a289cb..ae98f6526a 100644 --- a/crates/assistant2/src/ui.rs +++ b/crates/assistant2/src/ui.rs @@ -1,3 +1,5 @@ +mod chat_message_header; mod composer; +pub use chat_message_header::*; pub use composer::*; diff --git a/crates/assistant2/src/ui/chat_message_header.rs b/crates/assistant2/src/ui/chat_message_header.rs new file mode 100644 index 0000000000..d0732116c6 --- /dev/null +++ b/crates/assistant2/src/ui/chat_message_header.rs @@ -0,0 +1,57 @@ +use client::User; +use std::sync::Arc; +use ui::{prelude::*, Avatar}; + +pub enum UserOrAssistant { + User(Option>), + Assistant, +} + +#[derive(IntoElement)] +pub struct ChatMessageHeader { + player: UserOrAssistant, + contexts: Vec<()>, +} + +impl ChatMessageHeader { + pub fn new(player: UserOrAssistant) -> Self { + Self { + player, + contexts: Vec::new(), + } + } +} + +impl RenderOnce for ChatMessageHeader { + fn render(self, _cx: &mut WindowContext) -> impl IntoElement { + let (username, avatar_uri) = match self.player { + UserOrAssistant::Assistant => ( + "Assistant".into(), + Some("https://zed.dev/assistant_avatar.png".into()), + ), + UserOrAssistant::User(Some(user)) => { + (user.github_login.clone(), Some(user.avatar_uri.clone())) + } + UserOrAssistant::User(None) => ("You".into(), None), + }; + + h_flex() + .justify_between() + .child( + h_flex() + .gap_3() + .map(|this| { + let avatar_size = rems(20.0 / 16.0); + if let Some(avatar_uri) = avatar_uri { + this.child(Avatar::new(avatar_uri).size(avatar_size)) + } else { + this.child(div().size(avatar_size)) + } + }) + .child(Label::new(username).color(Color::Default)), + ) + .child(div().when(!self.contexts.is_empty(), |this| { + this.child(Label::new(self.contexts.len().to_string()).color(Color::Muted)) + })) + } +}