Add streaming context

Author: cpsievert

js/chat/chat.ts

@@ -77,11 +77,11 @@ class ChatMessage extends LightElement {
 
     return html`
       <div class="message-icon">${unsafeHTML(icon)}</div>
-      <shiny-markdown
+      <shiny-markdown-stream
         content=${this.content}
         content_type=${this.content_type}
         streaming=${this.streaming}
-      ></shiny-markdown>
+      ></shiny-markdown-stream>
     `;
   }
 }

js/markdown-stream/markdown-stream.ts

@@ -1,22 +1,35 @@
-import { LitElement,html } from "lit";
+import { LitElement, html } from "lit";
 import { unsafeHTML } from "lit-html/directives/unsafe-html.js";
 import { property } from "lit/decorators.js";
 
 import ClipboardJS from "clipboard";
 import { sanitize } from "dompurify";
 import hljs from "highlight.js/lib/common";
-import { Renderer,parse } from "marked";
+import { Renderer, parse } from "marked";
 
-import { LightElement,createElement,createSVGIcon } from "../utils/_utils";
+import { LightElement, createElement, createSVGIcon } from "../utils/_utils";
 
 type ContentType = "markdown" | "semi-markdown" | "html" | "text";
 
-type Message = {
+type ContentMessage = {
   id: string;
   content: string;
   operation: "append" | "replace";
 };
 
+type IsStreamingMessage = {
+  id: string;
+  isStreaming: boolean;
+};
+
+// Type guard
+function isStreamingMessage(
+  message: ContentMessage | IsStreamingMessage
+): message is IsStreamingMessage {
+  return "isStreaming" in message;
+}
+
+// SVG dot to indicate content is currently streaming
 const SVG_DOT = createSVGIcon(
   '<svg width="12" height="12" xmlns="http://www.w3.org/2000/svg" class="chat-streaming-dot" style="margin-left:.25em;margin-top:-.25em"><circle cx="6" cy="6" r="6"/></svg>'
 );
@@ -66,7 +79,7 @@ class MarkdownElement extends LightElement {
     if (changedProperties.has("content")) {
       this.#highlightAndCodeCopy();
       if (this.streaming) this.#appendStreamingDot();
-      // TODO: throw an event here that we're done and catch it in SHINY_CHAT_MESSAGE
+      // TODO: throw an event here that content has rendered and catch it in SHINY_CHAT_MESSAGE
       // requestScroll(this, this.streaming);
     }
     if (changedProperties.has("streaming")) {
@@ -110,28 +123,37 @@ class MarkdownElement extends LightElement {
   }
 }
 
-// TODO: is it a problem if this gets imported multiple times?
-customElements.define("shiny-markdown-stream", MarkdownElement);
+// TODO: it is a problem if this runs twice in the browser?
+if (!customElements.get("shiny-markdown-stream")) {
+  customElements.define("shiny-markdown-stream", MarkdownElement);
+}
+
+function handleMessage(message: ContentMessage | IsStreamingMessage): void {
+  const el = document.getElementById(message.id) as MarkdownElement;
+
+  if (!el) {
+    console.error(`Element with id ${message.id} not found`);
+    return;
+  }
+
+  if (isStreamingMessage(message)) {
+    el.streaming = message.isStreaming;
+    return;
+  }
+
+  if (message.operation === "replace") {
+    el.setAttribute("content", message.content);
+  } else if (message.operation === "append") {
+    const content = el.getAttribute("content");
+    el.setAttribute("content", content + message.content);
+  } else {
+    console.error(`Unknown operation: ${message.operation}`);
+  }
+}
 
+// TODO: it is a problem if this runs twice in the browser?
 $(function () {
-  Shiny.addCustomMessageHandler(
-    "shinyMarkdownStreamMessage",
-    function (message: Message) {
-      const el = document.getElementById(message.id);
-      if (!el) {
-        console.error(`Element with id ${message.id} not found`);
-        return;
-      }
-      if (message.operation === "replace") {
-        el.setAttribute("content", message.content);
-      } else if (message.operation === "append") {
-        const content = el.getAttribute("content");
-        el.setAttribute("content", content + message.content);
-      } else {
-        console.error(`Unknown operation: ${message.operation}`);
-      }
-    }
-  );
+  Shiny.addCustomMessageHandler("shinyMarkdownStreamMessage", handleMessage);
 });
 
-export { MarkdownElement,contentToHTML };
+export { MarkdownElement, contentToHTML };

shiny/ui/_markdown_stream.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from typing import Iterable, Literal
 
 from .. import reactive
@@ -15,12 +16,17 @@
 StreamingContentType = Literal["markdown", "semi-markdown", "html", "text"]
 
 
-class Message(TypedDict):
+class ContentMessage(TypedDict):
     id: str
     content: str
     operation: Literal["append", "replace"]
 
 
+class isStreamingMessage(TypedDict):
+    id: str
+    isStreaming: bool
+
+
 @add_example()
 class MarkdownStream:
     """
@@ -85,16 +91,17 @@ def stream(self, content: Iterable[str], clear: bool = True):
         async def _task():
             if clear:
                 self._replace("")
-            for c in content:
-                self._append(c)
+            with self._streaming_dot():
+                for c in content:
+                    self._append(c)
 
         _task()
 
     def update(self, content: str):
         self._replace(content)
 
     def _append(self, content: str):
-        msg: Message = {
+        msg: ContentMessage = {
             "id": self.id,
             "content": content,
             "operation": "append",
@@ -103,15 +110,32 @@ def _append(self, content: str):
         self._send_custom_message(msg)
 
     def _replace(self, content: str):
-        msg: Message = {
+        msg: ContentMessage = {
             "id": self.id,
             "content": content,
             "operation": "replace",
         }
 
         self._send_custom_message(msg)
 
-    def _send_custom_message(self, msg: Message):
+    @contextmanager
+    def _streaming_dot(self):
+        start: isStreamingMessage = {
+            "id": self.id,
+            "isStreaming": True,
+        }
+        self._send_custom_message(start)
+
+        try:
+            yield
+        finally:
+            end: isStreamingMessage = {
+                "id": self.id,
+                "isStreaming": False,
+            }
+            self._send_custom_message(end)
+
+    def _send_custom_message(self, msg: ContentMessage | isStreamingMessage):
         if self._session.is_stub_session():
             return
         self._session._send_message_sync(