Add .get_latest_stream_result() method to Chat and MarkdownStream

Author: cpsievert

shiny/ui/_chat.py

@@ -209,6 +209,10 @@ def __init__(
                 reactive.Value(None)
             )
 
+            self._latest_stream: reactive.Value[
+                reactive.ExtendedTask[[], str] | None
+            ] = reactive.Value(None)
+
             # TODO: deprecate messages once we start promoting managing LLM message
             # state through other means
             @reactive.effect
@@ -573,6 +577,8 @@ async def _stream_task():
 
         _stream_task()
 
+        self._latest_stream.set(_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
         @reactive.effect
@@ -584,6 +590,33 @@ async def _handle_error():
 
         return _stream_task
 
+    def get_latest_stream_result(self) -> str | None:
+        """
+        Reactively read the latest message stream result.
+
+        This method reads a reactive value containing the result of the latest
+        `.append_message_stream()`. Therefore, this method must be called in a reactive
+        context (e.g., a render function, a :func:`~shiny.reactive.calc`, or a
+        :func:`~shiny.reactive.effect`).
+
+        Returns
+        -------
+        :
+            The result of the latest stream (a string).
+
+        Raises
+        ------
+        :
+            A silent exception if no stream has completed yet.
+        """
+        stream = self._latest_stream()
+        if stream is None:
+            from .. import req
+
+            req(False)
+        else:
+            return stream.result()
+
     async def _append_message_stream(self, message: AsyncIterable[Any]):
         id = _utils.private_random_id()
 

shiny/ui/_markdown_stream.py

@@ -7,7 +7,7 @@
 from .._docstring import add_example
 from .._namespaces import resolve_id
 from .._typing_extensions import TypedDict
-from ..session import require_active_session
+from ..session import require_active_session, session_context
 from ..types import NotifyException
 from ..ui.css import CssUnit, as_css_unit
 from . import Tag
@@ -86,6 +86,11 @@ def __init__(
 
         self.on_error = on_error
 
+        with session_context(self._session):
+            self._latest_stream: reactive.Value[
+                reactive.ExtendedTask[[], str] | None
+            ] = reactive.Value(None)
+
     async def stream(
         self,
         content: Union[Iterable[str], AsyncIterable[str]],
@@ -146,6 +151,33 @@ async def _handle_error():
 
         return _task
 
+    def get_latest_stream_result(self) -> str | None:
+        """
+        Reactively read the latest stream result.
+
+        This method reads a reactive value containing the result of the latest
+        `.stream()`. Therefore, this method must be called in a reactive context (e.g.,
+        a render function, a :func:`~shiny.reactive.calc`, or a
+        :func:`~shiny.reactive.effect`).
+
+        Returns
+        -------
+        :
+            The result of the latest stream (a string).
+
+        Raises
+        ------
+        :
+            A silent exception if no stream has completed yet.
+        """
+        stream = self._latest_stream()
+        if stream is None:
+            from .. import req
+
+            req(False)
+        else:
+            return stream.result()
+
     async def clear(self):
         """
         Empty the UI element of the `MarkdownStream`.

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

@@ -1,9 +1,6 @@
 import asyncio
-from typing import Union
 
-from shiny import reactive
 from shiny.express import render, ui
-from shiny.reactive import ExtendedTask
 
 chat = ui.Chat("chat")
 
@@ -19,16 +16,9 @@ async def stream_generator():
 
 @chat.on_user_submit
 async def _(message: str):
-    stream = await chat.append_message_stream(stream_generator())
-    current_stream.set(stream)
-
-
-current_stream: reactive.value[Union[ExtendedTask[[], str], None]] = reactive.value(
-    None
-)
+    await chat.append_message_stream(stream_generator())
 
 
 @render.code
-def stream_result_ui():
-    stream = current_stream()
-    return stream.result() if stream else None
+async def stream_result_ui():
+    return chat.get_latest_stream_result()