Add the ability to change the icon when appending a message; rely less on DOM state for icon management

Author: cpsievert

js/chat/chat.scss

@@ -13,10 +13,6 @@ shiny-chat-container {
     margin-bottom: 0;
   }
 
-  [data-icon="assistant"] {
-    display: none;
-  }
-  
   .suggestion,
   [data-suggestion] {
     cursor: pointer;

js/chat/chat.ts

@@ -15,6 +15,7 @@ type Message = {
   role: "user" | "assistant";
   chunk_type: "message_start" | "message_end" | null;
   content_type: ContentType;
+  icon?: string;
   operation: "append" | null;
 };
 type ShinyChatMessage = {
@@ -60,20 +61,19 @@ class ChatMessage extends LightElement {
   @property() content = "...";
   @property() content_type: ContentType = "markdown";
   @property({ type: Boolean, reflect: true }) streaming = false;
+  @property() icon = "";
 
   render() {
-    const noContent = this.content.trim().length === 0;
-    const defaultIcon = noContent ? ICONS.dots_fade : ICONS.robot;
+    let msg_icon = ICONS.dots_fade;
 
-    // Check if there's an existing message-icon element
-    const userIcon = this.querySelector(".message-icon");
-    const icon =
-      userIcon && !noContent
-        ? userIcon
-        : html`<div class="message-icon">${unsafeHTML(defaultIcon)}</div>`;
+    // Show dots until we have content
+    const hasContent = this.content.trim().length !== 0;
+    if (hasContent) {
+      msg_icon = this.icon || ICONS.robot;
+    }
 
     return html`
-      ${icon}
+      <div class="message-icon">${unsafeHTML(msg_icon)}</div>
       <shiny-markdown-stream
         content=${this.content}
         content-type=${this.content_type}
@@ -269,6 +269,7 @@ class ChatInput extends LightElement {
 }
 
 class ChatContainer extends LightElement {
+  @property() icon = "";
 
   private get input(): ChatInput {
     return this.querySelector(CHAT_INPUT_TAG) as ChatInput;
@@ -355,17 +356,12 @@ class ChatContainer extends LightElement {
 
     const TAG_NAME =
       message.role === "user" ? CHAT_USER_MESSAGE_TAG : CHAT_MESSAGE_TAG;
-    const msg = createElement(TAG_NAME, message);
 
-    if (message.role !== "user") {
-      const iconTemplate = this.querySelector('div[data-icon="assistant"]');
-      if (iconTemplate) {
-        const icon = iconTemplate.cloneNode(true) as HTMLDivElement;
-        icon.className = "message-icon";
-        msg.appendChild(icon);
-      }
+    if (this.icon) {
+      message.icon = message.icon || this.icon;
     }
 
+    const msg = createElement(TAG_NAME, message);
     this.messages.appendChild(msg);
 
     if (finalize) {

shiny/ui/_chat.py

@@ -17,7 +17,7 @@
 )
 from weakref import WeakValueDictionary
 
-from htmltools import HTML, Tag, TagAttrValue, css, div
+from htmltools import HTML, Tag, TagAttrValue, css
 
 from .. import _utils, reactive
 from .._deprecated import warn_deprecated
@@ -489,7 +489,12 @@ def messages(
 
         return tuple(res)
 
-    async def append_message(self, message: Any) -> None:
+    async def append_message(
+        self,
+        message: Any,
+        *,
+        assistant_icon: HTML | Tag | None = None,
+    ):
         """
         Append a message to the chat.
 
@@ -502,6 +507,9 @@ async def append_message(self, message: Any) -> None:
             Content strings are interpreted as markdown and rendered to HTML on the
             client. Content may also include specially formatted **input suggestion**
             links (see note below).
+        assistant_icon
+            An optional icon to display next to the assistant message. The icon can be
+            any HTML element (e.g., an :func:`~shiny.ui.img` tag) or a string of HTML.
 
         Note
         ----
@@ -531,10 +539,15 @@ async def append_message(self, message: Any) -> None:
         similar) is specified in model's completion method.
         ```
         """
-        await self._append_message(message)
+        await self._append_message(message, icon=assistant_icon)
 
     async def _append_message(
-        self, message: Any, *, chunk: ChunkOption = False, stream_id: str | None = None
+        self,
+        message: Any,
+        *,
+        chunk: ChunkOption = False,
+        stream_id: str | None = None,
+        icon: HTML | Tag | None = None,
     ) -> None:
         # If currently we're in a stream, handle other messages (outside the stream) later
         if not self._can_append_message(stream_id):
@@ -564,9 +577,18 @@ async def _append_message(
         if msg is None:
             return
         self._store_message(msg, chunk=chunk)
-        await self._send_append_message(msg, chunk=chunk)
+        await self._send_append_message(
+            msg,
+            chunk=chunk,
+            icon=icon,
+        )
 
-    async def append_message_stream(self, message: Iterable[Any] | AsyncIterable[Any]):
+    async def append_message_stream(
+        self,
+        message: Iterable[Any] | AsyncIterable[Any],
+        *,
+        assistant_icon: HTML | Tag | None = None,
+    ):
         """
         Append a message as a stream of message chunks.
 
@@ -579,6 +601,9 @@ async def append_message_stream(self, message: Iterable[Any] | AsyncIterable[Any
             OpenAI, Anthropic, Ollama, and others. Content strings are interpreted as
             markdown and rendered to HTML on the client. Content may also include
             specially formatted **input suggestion** links (see note below).
+        assistant_icon
+            An optional icon to display next to the assistant message. The icon can be
+            any HTML element (e.g., an :func:`~shiny.ui.img` tag) or a string of HTML.
 
         Note
         ----
@@ -614,7 +639,7 @@ async def append_message_stream(self, message: Iterable[Any] | AsyncIterable[Any
         # Run the stream in the background to get non-blocking behavior
         @reactive.extended_task
         async def _stream_task():
-            await self._append_message_stream(message)
+            await self._append_message_stream(message, icon=assistant_icon)
 
         _stream_task()
 
@@ -627,11 +652,15 @@ async def _handle_error():
                 await self._raise_exception(e)
             _handle_error.destroy()  # type: ignore
 
-    async def _append_message_stream(self, message: AsyncIterable[Any]):
+    async def _append_message_stream(
+        self,
+        message: AsyncIterable[Any],
+        icon: HTML | Tag | None = None,
+    ):
         id = _utils.private_random_id()
 
         empty = ChatMessage(content="", role="assistant")
-        await self._append_message(empty, chunk="start", stream_id=id)
+        await self._append_message(empty, chunk="start", stream_id=id, icon=icon)
 
         try:
             async for msg in message:
@@ -659,6 +688,7 @@ async def _send_append_message(
         self,
         message: TransformedMessage,
         chunk: ChunkOption = False,
+        icon: HTML | Tag | None = None,
     ):
         if message["role"] == "system":
             # System messages are not displayed in the UI
@@ -678,13 +708,17 @@ async def _send_append_message(
         content = message["content_client"]
         content_type = "html" if isinstance(content, HTML) else "markdown"
 
+        # TODO: pass along dependencies for both content and icon (if any)
         msg = ClientMessage(
             content=str(content),
             role=message["role"],
             content_type=content_type,
             chunk_type=chunk_type,
         )
 
+        if icon is not None:
+            msg["icon"] = str(icon)
+
         # print(msg)
 
         await self._send_custom_message(msg_type, msg)
@@ -1187,34 +1221,26 @@ def chat_ui(
             raise ValueError("Each message must be a string or a dictionary.")
 
         if msg["role"] == "user":
-            msg_tag = Tag("shiny-user-message", content=msg["content"])
+            tag_name = "shiny-user-message"
         else:
-            msg_tag = Tag(
-                "shiny-chat-message",
-                (
-                    None
-                    if icon_assistant is None
-                    else div(icon_assistant, class_="message-icon")
-                ),
-                content=msg["content"],
-            )
+            tag_name = "shiny-chat-message"
 
-        message_tags.append(msg_tag)
+        message_tags.append(Tag(tag_name, content=msg["content"]))
+
+    html_deps = None
+    if isinstance(icon_assistant, Tag):
+        html_deps = icon_assistant.get_dependencies()
 
     res = Tag(
         "shiny-chat-container",
-        (
-            None
-            if icon_assistant is None
-            else div(HTML(icon_assistant), data_icon="assistant")
-        ),
         Tag("shiny-chat-messages", *message_tags),
         Tag(
             "shiny-chat-input",
             id=f"{id}_user_input",
             placeholder=placeholder,
         ),
         chat_deps(),
+        html_deps,
         {
             "style": css(
                 width=as_css_unit(width),
@@ -1224,6 +1250,7 @@ def chat_ui(
         id=id,
         placeholder=placeholder,
         fill=fill,
+        icon=str(icon_assistant) if icon_assistant else None,
         **kwargs,
     )
 

shiny/ui/_chat_types.py

@@ -4,6 +4,8 @@
 
 from htmltools import HTML
 
+from .._typing_extensions import NotRequired
+
 Role = Literal["assistant", "user", "system"]
 
 
@@ -27,3 +29,4 @@ class TransformedMessage(TypedDict):
 class ClientMessage(ChatMessage):
     content_type: Literal["markdown", "html"]
     chunk_type: Literal["message_start", "message_end"] | None
+    icon: NotRequired[str]