assistant2: Setup storybook (#11228)

This PR sets up the `assistant2` crate with the storybook so that UI
elements can be iterated on in isolation.

To start, we have some stories for the `ChatMessage` component:

```sh
cargo run -p storybook -- components/assistant_chat_message
```

<img width="1233" alt="Screenshot 2024-04-30 at 5 20 03 PM"
src="https://github.com/zed-industries/zed/assets/1486634/510967ea-0e9b-4fa9-94fb-421ee74bcc45">

Release Notes:

- N/A

Cargo.lock

@@ -401,6 +401,7 @@ dependencies = [
  "serde",
  "serde_json",
  "settings",
+ "story",
  "theme",
  "ui",
  "util",
@@ -9498,6 +9499,7 @@ name = "storybook"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "assistant2",
  "clap 4.4.4",
  "collab_ui",
  "ctrlc",

crates/assistant2/Cargo.toml

@@ -5,9 +5,16 @@ edition = "2021"
 publish = false
 license = "GPL-3.0-or-later"
 
+[lints]
+workspace = true
+
 [lib]
 path = "src/assistant2.rs"
 
+[features]
+default = []
+stories = ["dep:story"]
+
 [dependencies]
 anyhow.workspace = true
 assistant_tooling.workspace = true
@@ -29,6 +36,7 @@ semantic_index.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 settings.workspace = true
+story = { workspace = true, optional = true }
 theme.workspace = true
 ui.workspace = true
 util.workspace = true
@@ -49,6 +57,3 @@ settings = { workspace = true, features = ["test-support"] }
 theme = { workspace = true, features = ["test-support"] }
 util = { workspace = true, features = ["test-support"] }
 workspace = { workspace = true, features = ["test-support"] }
-
-[lints]
-workspace = true

crates/assistant2/src/assistant2.rs

@@ -1,7 +1,7 @@
 mod assistant_settings;
 mod completion_provider;
 mod tools;
-mod ui;
+pub mod ui;
 
 use ::ui::{div, prelude::*, Color, ViewContext};
 use anyhow::{Context, Result};
@@ -222,7 +222,7 @@ impl FocusableView for AssistantPanel {
     }
 }
 
-struct AssistantChat {
+pub struct AssistantChat {
     model: String,
     messages: Vec<ChatMessage>,
     list_state: ListState,
@@ -574,12 +574,15 @@ impl AssistantChat {
                 .map(|element| {
                     if self.editing_message_id.as_ref() == Some(id) {
                         element.child(Composer::new(
-                            cx.view().downgrade(),
-                            self.model.clone(),
                             body.clone(),
                             self.user_store.read(cx).current_user(),
                             self.can_submit(),
                             self.tool_registry.clone(),
+                            crate::ui::ModelSelector::new(
+                                cx.view().downgrade(),
+                                self.model.clone(),
+                            )
+                            .into_any_element(),
                         ))
                     } else {
                         element
@@ -744,18 +747,18 @@ impl Render for AssistantChat {
             .text_color(Color::Default.color(cx))
             .child(list(self.list_state.clone()).flex_1())
             .child(Composer::new(
-                cx.view().downgrade(),
-                self.model.clone(),
                 self.composer_editor.clone(),
                 self.user_store.read(cx).current_user(),
                 self.can_submit(),
                 self.tool_registry.clone(),
+                crate::ui::ModelSelector::new(cx.view().downgrade(), self.model.clone())
+                    .into_any_element(),
             ))
     }
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
-struct MessageId(usize);
+pub struct MessageId(usize);
 
 impl MessageId {
     fn post_inc(&mut self) -> Self {

crates/assistant2/src/ui.rs

@@ -1,5 +1,11 @@
 mod chat_message;
 mod composer;
 
+#[cfg(feature = "stories")]
+mod stories;
+
 pub use chat_message::*;
 pub use composer::*;
+
+#[cfg(feature = "stories")]
+pub use stories::*;

crates/assistant2/src/ui/composer.rs

@@ -1,7 +1,7 @@
 use assistant_tooling::ToolRegistry;
 use client::User;
 use editor::{Editor, EditorElement, EditorStyle};
-use gpui::{FontStyle, FontWeight, TextStyle, View, WeakView, WhiteSpace};
+use gpui::{AnyElement, FontStyle, FontWeight, TextStyle, View, WeakView, WhiteSpace};
 use settings::Settings;
 use std::sync::Arc;
 use theme::ThemeSettings;
@@ -11,30 +11,27 @@ use crate::{AssistantChat, CompletionProvider, Submit, SubmitMode};
 
 #[derive(IntoElement)]
 pub struct Composer {
-    assistant_chat: WeakView<AssistantChat>,
-    model: String,
     editor: View<Editor>,
     player: Option<Arc<User>>,
     can_submit: bool,
     tool_registry: Arc<ToolRegistry>,
+    model_selector: AnyElement,
 }
 
 impl Composer {
     pub fn new(
-        assistant_chat: WeakView<AssistantChat>,
-        model: String,
         editor: View<Editor>,
         player: Option<Arc<User>>,
         can_submit: bool,
         tool_registry: Arc<ToolRegistry>,
+        model_selector: AnyElement,
     ) -> Self {
         Self {
-            assistant_chat,
-            model,
             editor,
             player,
             can_submit,
             tool_registry,
+            model_selector,
         }
     }
 }
@@ -150,7 +147,7 @@ impl RenderOnce for Composer {
                         h_flex()
                             .w_full()
                             .justify_between()
-                            .child(ModelSelector::new(self.assistant_chat, self.model))
+                            .child(self.model_selector)
                             .children(self.tool_registry.status_views().iter().cloned()),
                     ),
             )
@@ -158,7 +155,7 @@ impl RenderOnce for Composer {
 }
 
 #[derive(IntoElement)]
-struct ModelSelector {
+pub struct ModelSelector {
     assistant_chat: WeakView<AssistantChat>,
     model: String,
 }

crates/assistant2/src/ui/stories.rs

@@ -0,0 +1,3 @@
+mod chat_message;
+
+pub use chat_message::*;

crates/assistant2/src/ui/stories/chat_message.rs

@@ -0,0 +1,99 @@
+use std::sync::Arc;
+
+use client::User;
+use story::{StoryContainer, StoryItem, StorySection};
+use ui::prelude::*;
+
+use crate::ui::{ChatMessage, UserOrAssistant};
+use crate::MessageId;
+
+pub struct ChatMessageStory;
+
+impl Render for ChatMessageStory {
+    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let user_1 = Arc::new(User {
+            id: 12345,
+            github_login: "iamnbutler".into(),
+            avatar_uri: "https://avatars.githubusercontent.com/u/1714999?v=4".into(),
+        });
+
+        StoryContainer::new(
+            "ChatMessage Story",
+            "crates/assistant2/src/ui/stories/chat_message.rs",
+        )
+        .child(
+            StorySection::new()
+                .child(StoryItem::new(
+                    "User chat message",
+                    ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::User(Some(user_1.clone())),
+                        Some(div().child("What can I do here?").into_any_element()),
+                        false,
+                        Box::new(|_, _| {}),
+                    ),
+                ))
+                .child(StoryItem::new(
+                    "User chat message (collapsed)",
+                    ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::User(Some(user_1.clone())),
+                        Some(div().child("What can I do here?").into_any_element()),
+                        true,
+                        Box::new(|_, _| {}),
+                    ),
+                )),
+        )
+        .child(
+            StorySection::new()
+                .child(StoryItem::new(
+                    "Assistant chat message",
+                    ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::Assistant,
+                        Some(div().child("You can talk to me!").into_any_element()),
+                        false,
+                        Box::new(|_, _| {}),
+                    ),
+                ))
+                .child(StoryItem::new(
+                    "Assistant chat message (collapsed)",
+                    ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::Assistant,
+                        Some(div().child("You can talk to me!").into_any_element()),
+                        true,
+                        Box::new(|_, _| {}),
+                    ),
+                )),
+        )
+        .child(
+            StorySection::new().child(StoryItem::new(
+                "Conversation between user and assistant",
+                v_flex()
+                    .gap_2()
+                    .child(ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::User(Some(user_1.clone())),
+                        Some(div().child("What is Rust??").into_any_element()),
+                        false,
+                        Box::new(|_, _| {}),
+                    ))
+                    .child(ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::Assistant,
+                        Some(div().child("Rust is a multi-paradigm programming language focused on performance and safety").into_any_element()),
+                        false,
+                        Box::new(|_, _| {}),
+                    ))
+                    .child(ChatMessage::new(
+                        MessageId(0),
+                        UserOrAssistant::User(Some(user_1)),
+                        Some(div().child("Sounds pretty cool!").into_any_element()),
+                        false,
+                        Box::new(|_, _| {}),
+                    )),
+            )),
+        )
+    }
+}

crates/storybook/Cargo.toml

@@ -14,6 +14,7 @@ path = "src/storybook.rs"
 
 [dependencies]
 anyhow.workspace = true
+assistant2 = { workspace = true, features = ["stories"] }
 clap = { workspace = true, features = ["derive", "string"] }
 collab_ui = { workspace = true, features = ["stories"] }
 ctrlc = "3.4"

crates/storybook/src/story_selector.rs

@@ -12,6 +12,7 @@ use ui::prelude::*;
 #[derive(Debug, PartialEq, Eq, Clone, Copy, strum::Display, EnumString, EnumIter)]
 #[strum(serialize_all = "snake_case")]
 pub enum ComponentStory {
+    AssistantChatMessage,
     AutoHeightEditor,
     Avatar,
     Button,
@@ -42,6 +43,9 @@ pub enum ComponentStory {
 impl ComponentStory {
     pub fn story(&self, cx: &mut WindowContext) -> AnyView {
         match self {
+            Self::AssistantChatMessage => {
+                cx.new_view(|_cx| assistant2::ui::ChatMessageStory).into()
+            }
             Self::AutoHeightEditor => AutoHeightEditorStory::new(cx).into(),
             Self::Avatar => cx.new_view(|_| ui::AvatarStory).into(),
             Self::Button => cx.new_view(|_| ui::ButtonStory).into(),