DEV: Improve `ai-streamer` API (#851)

In preparation for applying the streaming animation elsewhere, we want to better improve the organization of folder structure and methods used in the `ai-streamer`

assets/javascripts/discourse/connectors/topic-map-expanded-after/ai-summary-box.gjs

@@ -18,7 +18,8 @@ import I18n from "discourse-i18n";
 import DMenu from "float-kit/components/d-menu";
 import DTooltip from "float-kit/components/d-tooltip";
 import AiSummarySkeleton from "../../components/ai-summary-skeleton";
-import { streamSummaryText } from "../../lib/ai-streamer";
+import streamUpdaterText from "../../lib/ai-streamer/progress-handlers";
+import SummaryUpdater from "../../lib/ai-streamer/updaters/summary-updater";
 
 export default class AiSummaryBox extends Component {
   @service siteSettings;
@@ -34,7 +35,7 @@ export default class AiSummaryBox extends Component {
   @tracked canRegenerate = false;
   @tracked loading = false;
   @tracked isStreaming = false;
-  oldRaw = null; // used for comparison in SummaryUpdater in lib/ai-streamer
+  oldRaw = null; // used for comparison in SummaryUpdater in lib/ai-streamer/updaters
   finalSummary = null;
 
   get outdatedSummaryWarningText() {
@@ -150,7 +151,7 @@ export default class AiSummaryBox extends Component {
     this.loading = false;
 
     this.isStreaming = true;
-    streamSummaryText(topicSummary, this);
+    streamUpdaterText(SummaryUpdater, topicSummary, this);
 
     if (update.done) {
       this.isStreaming = false;
@@ -219,6 +220,7 @@ export default class AiSummaryBox extends Component {
               <article
                 class={{concatClass
                   "ai-summary-box"
+                  "streamable-content"
                   (if this.isStreaming "streaming")
                 }}
               >

assets/javascripts/discourse/lib/ai-streamer/progress-handlers.js

@@ -1,6 +1,5 @@
 import { later } from "@ember/runloop";
-import loadMorphlex from "discourse/lib/load-morphlex";
-import { cook } from "discourse/lib/text";
+import PostUpdater from "./updaters/post-updater";
 
 const PROGRESS_INTERVAL = 40;
 const GIVE_UP_INTERVAL = 60000;
@@ -9,6 +8,13 @@ const MAX_FLUSH_TIME = 800;
 
 let progressTimer = null;
 
+/**
+ * Finds the last non-empty child element or text node of a given DOM element.
+ * Iterates backward through the element's child nodes and skips over empty text nodes.
+ *
+ * @param {HTMLElement} element - The DOM element to inspect.
+ * @returns {Node} - The last non-empty child node or null if none found.
+ */
 function lastNonEmptyChild(element) {
   let lastChild = element.lastChild;
   while (
@@ -21,6 +27,13 @@ function lastNonEmptyChild(element) {
   return lastChild;
 }
 
+/**
+ * Adds a progress dot (a span element with a "progress-dot" class) at the end of the
+ * last non-empty block within a given DOM element. This is used to visually indicate
+ * progress while content is being streamed.
+ *
+ * @param {HTMLElement} element - The DOM element to which the progress dot will be added.
+ */
 export function addProgressDot(element) {
   let lastBlock = element;
 
@@ -42,143 +55,14 @@ export function addProgressDot(element) {
   lastBlock.appendChild(dotElement);
 }
 
-// this is the interface we need to implement
-// for a streaming updater
-class StreamUpdater {
-  set streaming(value) {
-    throw "not implemented";
-  }
-
-  async setCooked() {
-    throw "not implemented";
-  }
-
-  async setRaw() {
-    throw "not implemented";
-  }
-
-  get element() {
-    throw "not implemented";
-  }
-
-  get raw() {
-    throw "not implemented";
-  }
-}
-
-class PostUpdater extends StreamUpdater {
-  morphingOptions = {
-    beforeAttributeUpdated: (element, attributeName) => {
-      return !(element.tagName === "DETAILS" && attributeName === "open");
-    },
-  };
-
-  constructor(postStream, postId) {
-    super();
-    this.postStream = postStream;
-    this.postId = postId;
-    this.post = postStream.findLoadedPost(postId);
-
-    if (this.post) {
-      this.postElement = document.querySelector(
-        `#post_${this.post.post_number}`
-      );
-    }
-  }
-
-  get element() {
-    return this.postElement;
-  }
-
-  set streaming(value) {
-    if (this.postElement) {
-      if (value) {
-        this.postElement.classList.add("streaming");
-      } else {
-        this.postElement.classList.remove("streaming");
-      }
-    }
-  }
-
-  async setRaw(value, done) {
-    this.post.set("raw", value);
-    const cooked = await cook(value);
-
-    // resets animation
-    this.element.classList.remove("streaming");
-    void this.element.offsetWidth;
-    this.element.classList.add("streaming");
-
-    const cookedElement = document.createElement("div");
-    cookedElement.innerHTML = cooked;
-
-    if (!done) {
-      addProgressDot(cookedElement);
-    }
-
-    await this.setCooked(cookedElement.innerHTML);
-  }
-
-  async setCooked(value) {
-    this.post.set("cooked", value);
-
-    (await loadMorphlex()).morphInner(
-      this.postElement.querySelector(".cooked"),
-      `<div>${value}</div>`,
-      this.morphingOptions
-    );
-  }
-
-  get raw() {
-    return this.post.get("raw") || "";
-  }
-}
-
-export class SummaryUpdater extends StreamUpdater {
-  constructor(topicSummary, componentContext) {
-    super();
-    this.topicSummary = topicSummary;
-    this.componentContext = componentContext;
-
-    if (this.topicSummary) {
-      this.summaryBox = document.querySelector("article.ai-summary-box");
-    }
-  }
-
-  get element() {
-    return this.summaryBox;
-  }
-
-  set streaming(value) {
-    if (this.element) {
-      if (value) {
-        this.componentContext.isStreaming = true;
-      } else {
-        this.componentContext.isStreaming = false;
-      }
-    }
-  }
-
-  async setRaw(value, done) {
-    this.componentContext.oldRaw = value;
-    const cooked = await cook(value);
-
-    await this.setCooked(cooked);
-
-    if (done) {
-      this.componentContext.finalSummary = cooked;
-    }
-  }
-
-  async setCooked(value) {
-    this.componentContext.text = value;
-  }
-
-  get raw() {
-    return this.componentContext.oldRaw || "";
-  }
-}
-
+/**
+ * Applies progress to a streaming operation, updating the raw and cooked text,
+ * handling progress dots, and stopping streaming when complete.
+ *
+ * @param {Object} status - The current streaming status object.
+ * @param {Object} updater - An instance of a stream updater (e.g., PostUpdater or SummaryUpdater).
+ * @returns {Promise<boolean>} - Resolves to true if streaming is complete, otherwise false.
+ */
 export async function applyProgress(status, updater) {
   status.startTime = status.startTime || Date.now();
 
@@ -235,6 +119,14 @@ export async function applyProgress(status, updater) {
   return status.done;
 }
 
+/**
+ * Handles progress updates for a post stream by applying the streaming status of
+ * each post and updating its content accordingly. This function ensures that progress
+ * is tracked and handled for multiple posts simultaneously.
+ *
+ * @param {Object} postStream - The post stream object containing the posts to be updated.
+ * @returns {Promise<boolean>} - Resolves to true if polling should continue, otherwise false.
+ */
 async function handleProgress(postStream) {
   const status = postStream.aiStreamingStatus;
 
@@ -257,22 +149,12 @@ async function handleProgress(postStream) {
   return keepPolling;
 }
 
-export function streamSummaryText(topicSummary, context) {
-  const summaryUpdater = new SummaryUpdater(topicSummary, context);
-
-  if (!progressTimer) {
-    progressTimer = later(async () => {
-      await applyProgress(topicSummary, summaryUpdater);
-
-      progressTimer = null;
-
-      if (!topicSummary.done) {
-        await applyProgress(topicSummary, summaryUpdater);
-      }
-    }, PROGRESS_INTERVAL);
-  }
-}
-
+/**
+ * Ensures that progress for a post stream is being updated. It starts a progress timer
+ * if one is not already active, and continues polling for progress updates at regular intervals.
+ *
+ * @param {Object} postStream - The post stream object containing the posts to be updated.
+ */
 function ensureProgress(postStream) {
   if (!progressTimer) {
     progressTimer = later(async () => {
@@ -287,7 +169,14 @@ function ensureProgress(postStream) {
   }
 }
 
-export default function streamText(postStream, data) {
+/**
+ * Streams the raw text for a post by tracking its status and applying progress updates.
+ * If streaming is already in progress, this function ensures it continues to update the content.
+ *
+ * @param {Object} postStream - The post stream object containing the post to be updated.
+ * @param {Object} data - The data object containing raw and cooked content of the post.
+ */
+export function streamPostText(postStream, data) {
   if (data.noop) {
     return;
   }
@@ -297,3 +186,28 @@ export default function streamText(postStream, data) {
   status[data.post_id] = data;
   ensureProgress(postStream);
 }
+
+/**
+ * A generalized function to handle streaming of content using any specified updater class.
+ * It applies progress updates to the content (raw and cooked) based on the given data.
+ * Use this function to stream content for Glimmer components.
+ *
+ * @param {Function} updaterClass - The updater class to be used for streaming (e.g., PostUpdater, SummaryUpdater).
+ * @param {Object} data - The data object containing the content to be streamed.
+ * @param {Object} context - Additional context required for the updater (typically the context of the Ember component).
+ */
+export default function streamUpdaterText(updaterClass, data, context) {
+  const updaterInstance = new updaterClass(data, context);
+
+  if (!progressTimer) {
+    progressTimer = later(async () => {
+      await applyProgress(data, updaterInstance);
+
+      progressTimer = null;
+
+      if (!data.done) {
+        await applyProgress(data, updaterInstance);
+      }
+    }, PROGRESS_INTERVAL);
+  }
+}

assets/javascripts/discourse/lib/ai-streamer/updaters/post-updater.js

@@ -0,0 +1,72 @@
+import loadMorphlex from "discourse/lib/load-morphlex";
+import { cook } from "discourse/lib/text";
+import { addProgressDot } from "../progress-handlers";
+import StreamUpdater from "./stream-updater";
+
+export default class PostUpdater extends StreamUpdater {
+  morphingOptions = {
+    beforeAttributeUpdated: (element, attributeName) => {
+      return !(element.tagName === "DETAILS" && attributeName === "open");
+    },
+  };
+
+  constructor(postStream, postId) {
+    super();
+    this.postStream = postStream;
+    this.postId = postId;
+    this.post = postStream.findLoadedPost(postId);
+
+    if (this.post) {
+      this.postElement = document.querySelector(
+        `#post_${this.post.post_number}`
+      );
+    }
+  }
+
+  get element() {
+    return this.postElement;
+  }
+
+  set streaming(value) {
+    if (this.postElement) {
+      if (value) {
+        this.postElement.classList.add("streaming");
+      } else {
+        this.postElement.classList.remove("streaming");
+      }
+    }
+  }
+
+  async setRaw(value, done) {
+    this.post.set("raw", value);
+    const cooked = await cook(value);
+
+    // resets animation
+    this.element.classList.remove("streaming");
+    void this.element.offsetWidth;
+    this.element.classList.add("streaming");
+
+    const cookedElement = document.createElement("div");
+    cookedElement.innerHTML = cooked;
+
+    if (!done) {
+      addProgressDot(cookedElement);
+    }
+
+    await this.setCooked(cookedElement.innerHTML);
+  }
+
+  async setCooked(value) {
+    this.post.set("cooked", value);
+
+    (await loadMorphlex()).morphInner(
+      this.postElement.querySelector(".cooked"),
+      `<div>${value}</div>`,
+      this.morphingOptions
+    );
+  }
+
+  get raw() {
+    return this.post.get("raw") || "";
+  }
+}

assets/javascripts/discourse/lib/ai-streamer/updaters/stream-updater.js

@@ -0,0 +1,24 @@
+/**
+ * Interface needed to implement for a streaming updater
+ */
+export default class StreamUpdater {
+  set streaming(value) {
+    throw "not implemented";
+  }
+
+  async setCooked() {
+    throw "not implemented";
+  }
+
+  async setRaw() {
+    throw "not implemented";
+  }
+
+  get element() {
+    throw "not implemented";
+  }
+
+  get raw() {
+    throw "not implemented";
+  }
+}

assets/javascripts/discourse/lib/ai-streamer/updaters/summary-updater.js

@@ -0,0 +1,47 @@
+import { cook } from "discourse/lib/text";
+import StreamUpdater from "./stream-updater";
+
+export default class SummaryUpdater extends StreamUpdater {
+  constructor(topicSummary, componentContext) {
+    super();
+    this.topicSummary = topicSummary;
+    this.componentContext = componentContext;
+
+    if (this.topicSummary) {
+      this.summaryBox = document.querySelector("article.ai-summary-box");
+    }
+  }
+
+  get element() {
+    return this.summaryBox;
+  }
+
+  set streaming(value) {
+    if (this.element) {
+      if (value) {
+        this.componentContext.isStreaming = true;
+      } else {
+        this.componentContext.isStreaming = false;
+      }
+    }
+  }
+
+  async setRaw(value, done) {
+    this.componentContext.oldRaw = value;
+    const cooked = await cook(value);
+
+    await this.setCooked(cooked);
+
+    if (done) {
+      this.componentContext.finalSummary = cooked;
+    }
+  }
+
+  async setCooked(value) {
+    this.componentContext.text = value;
+  }
+
+  get raw() {
+    return this.componentContext.oldRaw || "";
+  }
+}

assets/javascripts/initializers/ai-bot-replies.js

@@ -5,7 +5,7 @@ import { withPluginApi } from "discourse/lib/plugin-api";
 import { registerWidgetShim } from "discourse/widgets/render-glimmer";
 import DebugAiModal from "../discourse/components/modal/debug-ai-modal";
 import ShareModal from "../discourse/components/modal/share-modal";
-import streamText from "../discourse/lib/ai-streamer";
+import { streamPostText } from "../discourse/lib/ai-streamer/progress-handlers";
 import copyConversation from "../discourse/lib/copy-conversation";
 const AUTO_COPY_THRESHOLD = 4;
 import AiBotHeaderIcon from "../discourse/components/ai-bot-header-icon";
@@ -50,7 +50,7 @@ function initializeAIBotReplies(api) {
     pluginId: "discourse-ai",
 
     onAIBotStreamedReply: function (data) {
-      streamText(this.model.postStream, data);
+      streamPostText(this.model.postStream, data);
     },
     subscribe: function () {
       this._super();

assets/stylesheets/common/streaming.scss

@@ -21,7 +21,7 @@
   color: var(--tertiary-medium);
 }
 
-.streaming .cooked p:last-child::after {
+.streamable-content.streaming .cooked p:last-child::after {
   @include progress-dot;
 }
 

test/javascripts/unit/lib/ai-streamer-test.js

@@ -3,7 +3,7 @@ import {
   addProgressDot,
   applyProgress,
   MIN_LETTERS_PER_INTERVAL,
-} from "discourse/plugins/discourse-ai/discourse/lib/ai-streamer";
+} from "discourse/plugins/discourse-ai/discourse/lib/ai-streamer/progress-handlers";
 
 class FakeStreamUpdater {
   constructor() {