feat: Configurable chat icon (#1853)

Co-authored-by: Carson <[email protected]>

Author: gadenbuie

CHANGELOG.md

@@ -31,6 +31,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * The `ui.Chat()` component's `.update_user_input()` method gains `submit` and `focus` options that allow you to submit the input on behalf of the user and to choose whether the input receives focus after the update. (#1851)
 
+* The assistant icons is now configurable via `ui.chat_ui()` (or the `ui.Chat.ui()` method in Shiny Express) or for individual messages in the `.append_message()` and `.append_message_stream()` methods of `ui.Chat()`. (#1853)
+
 ### Bug fixes
 
 * `ui.Chat()` now correctly handles new `ollama.chat()` return value introduced in `ollama` v0.4. (#1787)

js/chat/chat.scss

@@ -73,12 +73,30 @@ shiny-chat-message {
   .message-icon {
     border-radius: 50%;
     border: var(--shiny-chat-border);
+    height: 2rem;
+    width: 2rem;
+    display: grid;
+    place-items: center;
+    overflow: clip;
+
     > * {
-      margin: 0.5rem;
-      height: 20px;
-      width: 20px;
+      // images and avatars are full-bleed
+      height: 100%;
+      width: 100%;
+      margin: 0 !important;
+      object-fit: contain;
+    }
+
+    > svg,
+    > .icon,
+    > .fa,
+    > .bi {
+      // icons and svgs need some padding within the circle
+      max-height: 66%;
+      max-width: 66%;
     }
   }
+
   /* Vertically center the 2nd column (message content) */
   shiny-markdown-stream {
     align-self: center;

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,10 +61,12 @@ 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 icon = noContent ? ICONS.dots_fade : ICONS.robot;
+    // Show dots until we have content
+    const isEmpty = this.content.trim().length === 0;
+    const icon = isEmpty ? ICONS.dots_fade : this.icon || ICONS.robot;
 
     return html`
       <div class="message-icon">${unsafeHTML(icon)}</div>
@@ -262,6 +265,7 @@ class ChatInput extends LightElement {
 }
 
 class ChatContainer extends LightElement {
+  @property({ attribute: "icon-assistant" }) iconAssistant = "";
   inputSentinelObserver?: IntersectionObserver;
 
   private get input(): ChatInput {
@@ -381,6 +385,11 @@ class ChatContainer extends LightElement {
 
     const TAG_NAME =
       message.role === "user" ? CHAT_USER_MESSAGE_TAG : CHAT_MESSAGE_TAG;
+
+    if (this.iconAssistant) {
+      message.icon = message.icon || this.iconAssistant;
+    }
+
     const msg = createElement(TAG_NAME, message);
     this.messages.appendChild(msg);
 

shiny/ui/_chat.py

@@ -17,7 +17,7 @@
 )
 from weakref import WeakValueDictionary
 
-from htmltools import HTML, Tag, TagAttrValue, css
+from htmltools import HTML, Tag, TagAttrValue, TagList, css
 
 from .. import _utils, reactive
 from .._deprecated import warn_deprecated
@@ -493,7 +493,12 @@ def messages(
 
         return tuple(res)
 
-    async def append_message(self, message: Any) -> None:
+    async def append_message(
+        self,
+        message: Any,
+        *,
+        icon: HTML | Tag | TagList | None = None,
+    ):
         """
         Append a message to the chat.
 
@@ -506,10 +511,14 @@ 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).
+        icon
+            An optional icon to display next to the message, currently only used for
+            assistant messages. The icon can be any HTML element (e.g., an
+            :func:`~shiny.ui.img` tag) or a string of HTML.
 
         Note
         ----
-        ``{.callout-note title="Input suggestions"}
+        :::{.callout-note title="Input suggestions"}
         Input suggestions are special links that send text to the user input box when
         clicked (or accessed via keyboard). They can be created in the following ways:
 
@@ -528,17 +537,22 @@ async def append_message(self, message: Any) -> None:
 
         Note that a user may also opt-out of submitting a suggestion by holding the
         `Alt/Option` key while clicking the suggestion link.
-        ```
+        :::
 
-        ```{.callout-note title="Streamed messages"}
+        :::{.callout-note title="Streamed messages"}
         Use `.append_message_stream()` instead of this method when `stream=True` (or
         similar) is specified in model's completion method.
-        ```
+        :::
         """
-        await self._append_message(message)
+        await self._append_message(message, icon=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 | TagList | 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):
@@ -568,9 +582,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],
+        *,
+        icon: HTML | Tag | None = None,
+    ):
         """
         Append a message as a stream of message chunks.
 
@@ -583,6 +606,10 @@ 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).
+        icon
+            An optional icon to display next to the message, currently only used for
+            assistant messages. The icon can be any HTML element (e.g., an
+            :func:`~shiny.ui.img` tag) or a string of HTML.
 
         Note
         ----
@@ -625,7 +652,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():
-            return await self._append_message_stream(message)
+            return await self._append_message_stream(message, icon=icon)
 
         _stream_task()
 
@@ -669,11 +696,15 @@ def get_latest_stream_result(self) -> str | None:
         else:
             return stream.result()
 
-    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:
@@ -702,6 +733,7 @@ async def _send_append_message(
         self,
         message: TransformedMessage,
         chunk: ChunkOption = False,
+        icon: HTML | Tag | TagList | None = None,
     ):
         if message["role"] == "system":
             # System messages are not displayed in the UI
@@ -721,13 +753,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)
@@ -1118,7 +1154,6 @@ async def _send_custom_message(self, handler: str, obj: ClientMessage | None):
 
 @add_example(ex_dir="../templates/chat/starters/hello")
 class ChatExpress(Chat):
-
     def ui(
         self,
         *,
@@ -1127,6 +1162,7 @@ def ui(
         width: CssUnit = "min(680px, 100%)",
         height: CssUnit = "auto",
         fill: bool = True,
+        icon_assistant: HTML | Tag | TagList | None = None,
         **kwargs: TagAttrValue,
     ) -> Tag:
         """
@@ -1148,6 +1184,10 @@ def ui(
         fill
             Whether the chat should vertically take available space inside a fillable
             container.
+        icon_assistant
+            The icon to use for the assistant chat messages. Can be a HTML or a tag in
+            the form of :class:`~htmltools.HTML` or :class:`~htmltools.Tag`. If `None`,
+            a default robot icon is used.
         kwargs
             Additional attributes for the chat container element.
         """
@@ -1158,6 +1198,7 @@ def ui(
             width=width,
             height=height,
             fill=fill,
+            icon_assistant=icon_assistant,
             **kwargs,
         )
 
@@ -1171,6 +1212,7 @@ def chat_ui(
     width: CssUnit = "min(680px, 100%)",
     height: CssUnit = "auto",
     fill: bool = True,
+    icon_assistant: HTML | Tag | TagList | None = None,
     **kwargs: TagAttrValue,
 ) -> Tag:
     """
@@ -1199,6 +1241,10 @@ def chat_ui(
         The height of the chat container.
     fill
         Whether the chat should vertically take available space inside a fillable container.
+    icon_assistant
+            The icon to use for the assistant chat messages. Can be a HTML or a tag in
+            the form of :class:`~htmltools.HTML` or :class:`~htmltools.Tag`. If `None`,
+            a default robot icon is used.
     kwargs
         Additional attributes for the chat container element.
     """
@@ -1226,6 +1272,10 @@ def chat_ui(
 
         message_tags.append(Tag(tag_name, content=msg["content"]))
 
+    html_deps = None
+    if isinstance(icon_assistant, (Tag, TagList)):
+        html_deps = icon_assistant.get_dependencies()
+
     res = Tag(
         "shiny-chat-container",
         Tag("shiny-chat-messages", *message_tags),
@@ -1235,6 +1285,7 @@ def chat_ui(
             placeholder=placeholder,
         ),
         chat_deps(),
+        html_deps,
         {
             "style": css(
                 width=as_css_unit(width),
@@ -1244,6 +1295,7 @@ def chat_ui(
         id=id,
         placeholder=placeholder,
         fill=fill,
+        icon_assistant=str(icon_assistant) if icon_assistant is not None 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]