Yield a MessageStream() instance with a .append() and .restore() method

Author: cpsievert

shiny/ui/_chat.py

@@ -590,12 +590,7 @@ async def append_message(
             icon=icon,
         )
 
-    async def append_message_chunk(
-        self,
-        message_chunk: Any,
-        *,
-        operation: Literal["append", "replace"] = "append",
-    ):
+    async def append_message_chunk(self, message_chunk: Any):
         """
         Append a message chunk to the current message stream.
 
@@ -606,8 +601,6 @@ async def append_message_chunk(
         ----------
         message_chunk
             A message chunk to inject.
-        operation
-            Whether to append or replace the *current* message stream content.
 
         Note
         ----
@@ -633,7 +626,6 @@ async def append_message_chunk(
         return await self._append_message_chunk(
             message_chunk,
             stream_id=stream_id,
-            operation=operation,
         )
 
     @asynccontextmanager
@@ -644,6 +636,12 @@ async def message_stream(self):
         A context manager for streaming messages into the chat. Note this stream
         can occur within a longer running `.append_message_stream()` or used on its own.
 
+        Yields
+        ------
+        :
+            A `MessageStream` instance with a method for `.append()`ing message chunks
+            and a method for `.restore()`ing the stream back to it's initial state.
+
         Note
         ----
         A useful pattern for displaying tool calls in a chat interface is for the
@@ -658,16 +656,14 @@ async def message_stream(self):
         self._message_stream_checkpoint = self._current_stream_message
 
         # No stream currently exists, start one
-        is_root_stream = not self._current_stream_id
+        stream_id = self._current_stream_id
+        is_root_stream = stream_id is None
         if is_root_stream:
-            await self._append_message_chunk(
-                "",
-                chunk="start",
-                stream_id=_utils.private_random_id(),
-            )
+            stream_id = _utils.private_random_id()
+            await self._append_message_chunk("", chunk="start", stream_id=stream_id)
 
         try:
-            yield
+            yield MessageStream(self, stream_id)
         finally:
             # Restore the previous stream state
             self._message_stream_checkpoint = old_checkpoint
@@ -677,7 +673,7 @@ async def message_stream(self):
                 await self._append_message_chunk(
                     "",
                     chunk="end",
-                    stream_id=cast(str, self._current_stream_id),
+                    stream_id=stream_id,
                 )
 
     async def _append_message_chunk(
@@ -1496,4 +1492,36 @@ def chat_ui(
     return res
 
 
+class MessageStream:
+    """"""
+
+    def __init__(self, chat: Chat, stream_id: str):
+        self._chat = chat
+        self._stream_id = stream_id
+
+    async def restore(self):
+        """
+        Restore the stream back to its initial state.
+        """
+        await self._chat._append_message_chunk(
+            "",
+            operation="replace",
+            stream_id=self._stream_id,
+        )
+
+    async def append(self, message_chunk: Any):
+        """
+        Append a message chunk to the stream.
+
+        Parameters
+        -----------
+        message_chunk
+            A message chunk to append to this stream
+        """
+        await self._chat._append_message_chunk(
+            message_chunk,
+            stream_id=self._stream_id,
+        )
+
+
 CHAT_INSTANCES: WeakValueDictionary[str, Chat] = WeakValueDictionary()

tests/playwright/shiny/components/chat/message-stream/app.py

@@ -3,7 +3,7 @@
 from shiny import reactive
 from shiny.express import input, render, ui
 
-SLEEP_TIME = 0.75
+SLEEP_TIME = 0.25
 
 ui.page_opts(title="Hello chat message streams")
 
@@ -37,11 +37,12 @@ async def _():
         " Progress: 50%",
         " Progress: 100%",
     ]
-    async with chat.message_stream():
+    async with chat.message_stream() as stream:
         for chunk in chunks:
-            await chat.append_message_chunk(chunk)
+            await stream.append(chunk)
             await asyncio.sleep(SLEEP_TIME)
-        await chat.append_message_chunk("Completed stream 1 ✅", operation="replace")
+        await stream.restore()
+        await stream.append("Completed stream 1 ✅")
 
 
 # TODO: add test here for nested .message_stream()
@@ -68,8 +69,12 @@ async def mock_tool():
         " Progress: 50%",
         " Progress: 100%",
     ]
-    for chunk in chunks:
-        await chat.append_message_chunk(chunk, operation="replace")
+    async with chat.message_stream() as stream:
+        for chunk in chunks:
+            await stream.append(chunk)
+            await asyncio.sleep(SLEEP_TIME)
+        await stream.restore()
+        await stream.append("Completed inner stream ✅")
 
 
 # TODO: more tests, like submitting input, etc.