refactor: Use shadow root and slots in ChatContainer and ChatMessage

Currently used to manage the assistant icon,
but this pattern can be used for other icons as well.

Author: gadenbuie

js/chat/chat.scss

@@ -57,32 +57,16 @@ shiny-chat-messages {
 }
 
 shiny-chat-message {
-  display: grid;
-  grid-template-columns: auto minmax(0, 1fr);
-  gap: 1rem;
   > * {
     height: fit-content;
   }
-  .message-icon {
-    border-radius: 50%;
-    border: var(--shiny-chat-border);
-    height: 2rem;
-    width: 2rem;
-    display: grid;
-    place-items: center;
-
-    > * {
-      height: 20px;
-      width: 20px;
-      margin: 0 !important;
-      max-height: 85%;
-      max-width: 85%;
-      object-fit: contain;
-    }
-  }
-  /* Vertically center the 2nd column (message content) */
-  shiny-markdown-stream {
-    align-self: center;
+  [slot="icon"] {
+    height: 20px;
+    width: 20px;
+    margin: 0 !important;
+    max-height: 85%;
+    max-width: 85%;
+    object-fit: contain;
   }
 }
 

js/chat/chat.ts

@@ -1,4 +1,4 @@
-import { LitElement, html } from "lit";
+import { LitElement, css, html } from "lit";
 import { unsafeHTML } from "lit-html/directives/unsafe-html.js";
 import { property } from "lit/decorators.js";
 
@@ -57,23 +57,63 @@ const ICONS = {
     '<svg width="24" height="24" fill="currentColor" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><style>.spinner_S1WN{animation:spinner_MGfb .8s linear infinite;animation-delay:-.8s}.spinner_Km9P{animation-delay:-.65s}.spinner_JApP{animation-delay:-.5s}@keyframes spinner_MGfb{93.75%,100%{opacity:.2}}</style><circle class="spinner_S1WN" cx="4" cy="12" r="3"/><circle class="spinner_S1WN spinner_Km9P" cx="12" cy="12" r="3"/><circle class="spinner_S1WN spinner_JApP" cx="20" cy="12" r="3"/></svg>',
 };
 
-class ChatMessage extends LightElement {
+class ChatMessage extends LitElement {
   @property() content = "...";
   @property() content_type: ContentType = "markdown";
   @property({ type: Boolean, reflect: true }) streaming = false;
-  @property() icon = "";
 
-  render() {
-    let msg_icon = ICONS.dots_fade;
+  static styles = css`
+    :host {
+      display: grid;
+      grid-template-columns: auto minmax(0, 1fr);
+      gap: 1rem;
+      align-items: start;
+    }
 
-    // Show dots until we have content
-    const hasContent = this.content.trim().length !== 0;
-    if (hasContent) {
-      msg_icon = this.icon || ICONS.robot;
+    :host > * {
+      height: fit-content;
+    }
+
+    .message-icon {
+      border-radius: 50%;
+      border: var(--shiny-chat-border);
+      height: 2rem;
+      width: 2rem;
+      display: grid;
+      place-items: center;
     }
 
+    slot[name="icon"] > * {
+      height: 20px;
+      width: 20px;
+      margin: 0 !important;
+      max-height: 85%;
+      max-width: 85%;
+      object-fit: contain;
+    }
+
+    /* Vertically center the 2nd column (message content) */
+    shiny-markdown-stream {
+      align-self: center;
+
+      p:first-child {
+        margin-block-start: 0;
+      }
+      p:last-child {
+        margin-block-end: 0;
+      }
+    }
+  `;
+
+  render() {
+    // Show dots until we have content
+    const hasContent = this.content.trim().length > 0;
+    const defaultIcon = hasContent ? ICONS.robot : ICONS.dots_fade;
+
     return html`
-      <div class="message-icon">${unsafeHTML(msg_icon)}</div>
+      <div class="message-icon">
+        <slot name="icon">${unsafeHTML(defaultIcon)}</slot>
+      </div>
       <shiny-markdown-stream
         content=${this.content}
         content-type=${this.content_type}
@@ -268,9 +308,7 @@ class ChatInput extends LightElement {
   }
 }
 
-class ChatContainer extends LightElement {
-  @property() icon = "";
-
+class ChatContainer extends LitElement {
   private get input(): ChatInput {
     return this.querySelector(CHAT_INPUT_TAG) as ChatInput;
   }
@@ -284,8 +322,30 @@ class ChatContainer extends LightElement {
     return last ? (last as ChatMessage) : null;
   }
 
+  private get iconAssistant(): Element | void {
+    const slot = this.shadowRoot?.querySelector(
+      'slot[name="icon-assistant"]'
+    ) as HTMLSlotElement;
+
+    if (!slot) return;
+
+    let icon: Element | undefined | null = slot.assignedElements()[0];
+    if (!icon) return;
+
+    if (icon?.matches(".icon-container")) {
+      // From Python/R we use a wrapper element because users may give raw HTML
+      icon = icon.firstElementChild;
+    }
+
+    return icon ? icon : undefined;
+  }
+
   render() {
-    return html``;
+    return html`
+      <slot name="icon-assistant" style="display: none"></slot>
+      <slot name="messages"></slot>
+      <slot name="input"></slot>
+    `;
   }
 
   firstUpdated(): void {
@@ -351,17 +411,42 @@ class ChatContainer extends LightElement {
     }
   }
 
+  #messageIcon(message: Message): HTMLElement | undefined {
+    if (message.role === "user") return;
+
+    let icon: HTMLElement | undefined;
+    if (message.icon) {
+      icon = document.createElement("div");
+      icon.innerHTML = message.icon;
+      if (icon.firstChild) {
+        icon = icon.firstChild.cloneNode() as HTMLElement;
+      } else {
+        icon = undefined;
+      }
+    }
+
+    if (!icon && this.iconAssistant) {
+      icon = this.iconAssistant.cloneNode(true) as HTMLElement;
+    }
+
+    if (!icon) return;
+    icon.setAttribute("slot", "icon");
+    return icon;
+  }
+
   #appendMessage(message: Message, finalize = true): void {
     this.#initMessage();
 
     const TAG_NAME =
       message.role === "user" ? CHAT_USER_MESSAGE_TAG : CHAT_MESSAGE_TAG;
 
-    if (this.icon) {
-      message.icon = message.icon || this.icon;
+    const msg = createElement(TAG_NAME, message);
+
+    const icon = this.#messageIcon(message);
+    if (icon) {
+      msg.appendChild(icon);
     }
 
-    const msg = createElement(TAG_NAME, message);
     this.messages.appendChild(msg);
 
     if (finalize) {

shiny/ui/_chat.py

@@ -1233,11 +1233,13 @@ def chat_ui(
 
     res = Tag(
         "shiny-chat-container",
-        Tag("shiny-chat-messages", *message_tags),
+        Tag("div", icon_assistant, class_="icon-container", slot="icon-assistant"),
+        Tag("shiny-chat-messages", *message_tags, slot="messages"),
         Tag(
             "shiny-chat-input",
             id=f"{id}_user_input",
             placeholder=placeholder,
+            slot="input",
         ),
         chat_deps(),
         html_deps,
@@ -1250,7 +1252,6 @@ def chat_ui(
         id=id,
         placeholder=placeholder,
         fill=fill,
-        icon=str(icon_assistant) if icon_assistant else None,
         **kwargs,
     )