Allow @chat.transform_assistant_response function to return None

CHANGELOG.md

@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * A handful of fixes for `ui.Chat()`, including:
   * A fix for use inside Shiny modules. (#1582)
   * `.messages(format="google")` now returns the correct role. (#1622)
+  * `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)
 

js/chat/chat.ts

@@ -275,7 +275,10 @@ class ChatContainer extends LightElement {
       this.#onAppendChunk
     );
     this.addEventListener("shiny-chat-clear-messages", this.#onClear);
-    this.addEventListener("shiny-chat-update-user-input", this.#onUpdateUserInput);
+    this.addEventListener(
+      "shiny-chat-update-user-input",
+      this.#onUpdateUserInput
+    );
     this.addEventListener(
       "shiny-chat-remove-loading-message",
       this.#onRemoveLoadingMessage
@@ -369,6 +372,7 @@ class ChatContainer extends LightElement {
 
     if (message.chunk_type === "message_end") {
       lastMessage.removeAttribute("is_streaming");
+      lastMessage.setAttribute("content", message.content);
       this.#finalizeMessage();
       return;
     }

shiny/ui/_chat.py

@@ -52,11 +52,11 @@
 # user input content types
 TransformUserInput = Callable[[str], Union[str, None]]
 TransformUserInputAsync = Callable[[str], Awaitable[Union[str, None]]]
-TransformAssistantResponse = Callable[[str], Union[str, HTML]]
-TransformAssistantResponseAsync = Callable[[str], Awaitable[Union[str, HTML]]]
-TransformAssistantResponseChunk = Callable[[str, str, bool], Union[str, HTML]]
+TransformAssistantResponse = Callable[[str], Union[str, HTML, None]]
+TransformAssistantResponseAsync = Callable[[str], Awaitable[Union[str, HTML, None]]]
+TransformAssistantResponseChunk = Callable[[str, str, bool], Union[str, HTML, None]]
 TransformAssistantResponseChunkAsync = Callable[
-    [str, str, bool], Awaitable[Union[str, HTML]]
+    [str, str, bool], Awaitable[Union[str, HTML, None]]
 ]
 TransformAssistantResponseFunction = Union[
     TransformAssistantResponse,
@@ -711,11 +711,11 @@ def transform_assistant_response(
         Parameters
         ----------
         fn
-            A function that takes a string and returns a string or
-            :class:`shiny.ui.HTML`. If `fn` returns a string, it gets interpreted and
-            parsed as a markdown on the client (and the resulting HTML is then
-            sanitized). If `fn` returns :class:`shiny.ui.HTML`, it will be displayed
-            as-is.
+            A function that takes a string and returns either a string,
+            :class:`shiny.ui.HTML`, or `None`. If `fn` returns a string, it gets
+            interpreted and parsed as a markdown on the client (and the resulting HTML
+            is then sanitized). If `fn` returns :class:`shiny.ui.HTML`, it will be
+            displayed as-is. If `fn` returns `None`, the response is effectively ignored.
 
         Note
         ----
@@ -774,16 +774,20 @@ async def _transform_message(
 
         if message["role"] == "user" and self._transform_user is not None:
             content = await self._transform_user(message["content"])
-            if content is None:
-                return None
-            res[key] = content
 
         elif message["role"] == "assistant" and self._transform_assistant is not None:
-            res[key] = await self._transform_assistant(
+            content = await self._transform_assistant(
                 message["content"],
                 chunk_content or "",
                 chunk == "end" or chunk is False,
             )
+        else:
+            return res
+
+        if content is None:
+            return None
+
+        res[key] = content
 
         return res
 

tests/playwright/shiny/components/chat/transform_assistant_stream/app.py

@@ -0,0 +1,26 @@
+from shiny import render
+from shiny.express import ui
+
+chat = ui.Chat(id="chat")
+chat.ui()
+
+
+@chat.transform_assistant_response
+def transform(content: str, chunk: str, done: bool):
+    if done:
+        return content + "...DONE!"
+    else:
+        return content
+
+
+@chat.on_user_submit
+async def _():
+    await chat.append_message_stream(("Simple ", "response"))
+
+
+"Message state:"
+
+
+@render.code
+def message_state():
+    return str(chat.messages())

tests/playwright/shiny/components/chat/transform_assistant_stream/test_chat_transform_assistant_stream.py

@@ -0,0 +1,31 @@
+from playwright.sync_api import Page, expect
+from utils.deploy_utils import skip_on_webkit
+
+from shiny.playwright import controller
+from shiny.run import ShinyAppProc
+
+
+@skip_on_webkit
+def test_validate_chat_transform_assistant(page: Page, local_app: ShinyAppProc) -> None:
+    page.goto(local_app.url)
+
+    chat = controller.Chat(page, "chat")
+    message_state = controller.OutputCode(page, "message_state")
+
+    # Wait for app to load
+    message_state.expect_value("()", timeout=30 * 1000)
+
+    expect(chat.loc).to_be_visible(timeout=30 * 1000)
+    expect(chat.loc_input_button).to_be_disabled()
+
+    chat.set_user_input("foo")
+    chat.send_user_input()
+    chat.expect_latest_message("Simple response...DONE!", timeout=30 * 1000)
+
+    message_state_expected = tuple(
+        [
+            {"content": "foo", "role": "user"},
+            {"content": "Simple response", "role": "assistant"},
+        ]
+    )
+    message_state.expect_value(str(message_state_expected))