Add messages parameter to chat_ui() and Chat.ui(); revert #1593 (…

…#1736)
posit-dev · Oct 15, 2024 · d0354b9 · d0354b9
1 parent 023de32
commit d0354b9
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 45 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -29,6 +29,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * Added [narwhals](https://posit-dev.github.io/py-narwhals) support for `@render.table`. This allows for any eager data frame supported by narwhals to be returned from a `@render.table` output method. (#1570)
 
+* `chat_ui()` and `Chat.ui()` gain a `messages` parameter for providing starting messages. (#1736)
+
 ### Other changes
 
 * Incorporated `orjson` for faster data serialization in `@render.data_frame` outputs. (#1570)
@@ -100,7 +102,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * A few fixes for `ui.Chat()`, including:
   * A fix for use inside Shiny modules. (#1582)
   * `.messages(format="google")` now returns the correct role. (#1622)
-  * `ui.Chat(messages)` are no longer dropped when dynamically rendered. (#1593)
   * `transform_assistant_response` can now return `None` and correctly handles change of content on the last chunk. (#1641)
 
 * An empty `ui.input_date()` value no longer crashes Shiny. (#1528)

diff --git a/js/chat/chat.ts b/js/chat/chat.ts
@@ -313,14 +313,7 @@ class ChatContainer extends LightElement {
   private resizeObserver!: ResizeObserver;
 
   render(): ReturnType<LitElement["render"]> {
-    const input_id = this.id + "_user_input";
-    return html`
-      <shiny-chat-messages></shiny-chat-messages>
-      <shiny-chat-input
-        id=${input_id}
-        placeholder=${this.placeholder}
-      ></shiny-chat-input>
-    `;
+    return html``;
   }
 
   firstUpdated(): void {

diff --git a/shiny/ui/_chat.py b/shiny/ui/_chat.py
@@ -8,6 +8,7 @@
     Callable,
     Iterable,
     Literal,
+    Optional,
     Sequence,
     Tuple,
     Union,
@@ -193,7 +194,8 @@ def __init__(
                 reactive.Value(None)
             )
 
-            # Initialize the chat with the provided messages
+            # TODO: deprecate messages once we start promoting managing LLM message
+            # state through other means
             @reactive.effect
             async def _init_chat():
                 for msg in messages:
@@ -233,6 +235,7 @@ async def _on_user_input():
     def ui(
         self,
         *,
+        messages: Optional[Sequence[str | ChatMessage]] = None,
         placeholder: str = "Enter a message...",
         width: CssUnit = "min(680px, 100%)",
         height: CssUnit = "auto",
@@ -247,6 +250,11 @@ def ui(
 
         Parameters
         ----------
+        messages
+            A sequence of messages to display in the chat. Each message can be either a
+            string or a dictionary with `content` and `role` keys. The `content` key
+            should contain the message text, and the `role` key can be "assistant" or
+            "user".
         placeholder
             Placeholder text for the chat input.
         width
@@ -261,6 +269,7 @@ def ui(
         """
         return chat_ui(
             id=self.id,
+            messages=messages,
             placeholder=placeholder,
             width=width,
             height=height,
@@ -572,7 +581,7 @@ async def append_message_stream(self, message: Iterable[Any] | AsyncIterable[Any
         async def _stream_task():
             await self._append_message_stream(message)
 
-        self._session.on_flushed(_stream_task, once=True)
+        _stream_task()
 
         # Since the task runs in the background (outside/beyond the current context,
         # if any), we need to manually raise any exceptions that occur
@@ -643,9 +652,7 @@ async def _send_append_message(
 
         # print(msg)
 
-        # When streaming (i.e., chunk is truthy), we can send messages immediately
-        # since we already waited for the flush in order to start the stream
-        await self._send_custom_message(msg_type, msg, on_flushed=chunk is False)
+        await self._send_custom_message(msg_type, msg)
         # TODO: Joe said it's a good idea to yield here, but I'm not sure why?
         # await asyncio.sleep(0)
 
@@ -1012,29 +1019,22 @@ def destroy(self):
     async def _remove_loading_message(self):
         await self._send_custom_message("shiny-chat-remove-loading-message", None)
 
-    async def _send_custom_message(
-        self, handler: str, obj: ClientMessage | None, on_flushed: bool = True
-    ):
-        async def _do_send():
-            await self._session.send_custom_message(
-                "shinyChatMessage",
-                {
-                    "id": self.id,
-                    "handler": handler,
-                    "obj": obj,
-                },
-            )
-
-        if on_flushed:
-            self._session.on_flushed(_do_send, once=True)
-        else:
-            await _do_send()
+    async def _send_custom_message(self, handler: str, obj: ClientMessage | None):
+        await self._session.send_custom_message(
+            "shinyChatMessage",
+            {
+                "id": self.id,
+                "handler": handler,
+                "obj": obj,
+            },
+        )
 
 
 @add_example(ex_dir="../api-examples/chat")
 def chat_ui(
     id: str,
     *,
+    messages: Optional[Sequence[str | ChatMessage]] = None,
     placeholder: str = "Enter a message...",
     width: CssUnit = "min(680px, 100%)",
     height: CssUnit = "auto",
@@ -1052,6 +1052,10 @@ def chat_ui(
     ----------
     id
         A unique identifier for the chat UI.
+    messages
+        A sequence of messages to display in the chat. Each message can be either a string
+        or a dictionary with a `content` and `role` key. The `content` key should contain
+        the message text, and the `role` key can be "assistant" or "user".
     placeholder
         Placeholder text for the chat input.
     width
@@ -1066,8 +1070,32 @@ def chat_ui(
 
     id = resolve_id(id)
 
+    message_tags: list[Tag] = []
+    if messages is None:
+        messages = []
+    for msg in messages:
+        if isinstance(msg, str):
+            msg = {"content": msg, "role": "assistant"}
+        elif isinstance(msg, dict):
+            if "content" not in msg:
+                raise ValueError("Each message must have a 'content' key.")
+            if "role" not in msg:
+                raise ValueError("Each message must have a 'role' key.")
+        else:
+            raise ValueError("Each message must be a string or a dictionary.")
+
+        message_tags.append(
+            Tag("shiny-chat-message", content=msg["content"], role=msg["role"])
+        )
+
     res = Tag(
         "shiny-chat-container",
+        Tag("shiny-chat-messages", *message_tags),
+        Tag(
+            "shiny-chat-input",
+            id=f"{id}_user_input",
+            placeholder=placeholder,
+        ),
         chat_deps(),
         {
             "style": css(

diff --git a/shiny/www/py-shiny/chat/chat.js b/shiny/www/py-shiny/chat/chat.js
diff --git a/shiny/www/py-shiny/chat/chat.js.map b/shiny/www/py-shiny/chat/chat.js.map
diff --git a/tests/playwright/shiny/components/chat/dynamic_ui/app.py b/tests/playwright/shiny/components/chat/dynamic_ui/app.py
@@ -1,8 +1,8 @@
 from shiny.express import render, ui
 
-chat = ui.Chat(id="chat", messages=["A starting message"])
+chat = ui.Chat(id="chat")
 
 
 @render.ui
 def chat_output():
-    return chat.ui()
+    return chat.ui(messages=["A starting message"])
diff --git a/tests/playwright/shiny/components/chat/stream/test_chat_stream.py b/tests/playwright/shiny/components/chat/stream/test_chat_stream.py
@@ -21,10 +21,10 @@ def test_validate_chat(page: Page, local_app: ShinyAppProc) -> None:
     expect(chat.loc_input_button).to_be_disabled()
 
     messages = [
-        "SECOND SECOND SECOND",
-        "FOURTH FOURTH FOURTH",
         "FIRST FIRST FIRST",
+        "SECOND SECOND SECOND",
         "THIRD THIRD THIRD",
+        "FOURTH FOURTH FOURTH",
         "FIFTH FIFTH FIFTH",
     ]
     # Allow for any whitespace between messages