feat: support new streaming mechanism in the playground (#601)

Yonom · Jul 28, 2024 · 633d6fe · 633d6fe
1 parent bbfdbec
commit 633d6fe
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/packages/react-playground/package.json b/packages/react-playground/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@assistant-ui/react-playground",
-  "version": "0.0.11",
+  "version": "0.0.12",
   "license": "MIT",
   "exports": {
     ".": {

diff --git a/packages/react-playground/src/lib/playground-runtime.ts b/packages/react-playground/src/lib/playground-runtime.ts
@@ -193,21 +193,31 @@ export class PlaygroundThreadRuntime implements ReactThreadRuntime {
     };
 
     try {
-      const result = await this.adapter.run({
+      const promiseOrGenerator = this.adapter.run({
         messages,
         abortSignal: this.abortController.signal,
         config: this.configProvider.getModelConfig(),
         onUpdate: updateMessage,
       });
-      if (result.status?.type === "running")
-        throw new Error(
-          "Unexpected running status returned from ChatModelAdapter",
-        );
+
+      if (Symbol.asyncIterator in promiseOrGenerator) {
+        for await (const r of promiseOrGenerator) {
+          updateMessage(r);
+        }
+      } else {
+        updateMessage(await promiseOrGenerator);
+      }
+
       this.abortController = null;
-      updateMessage({
-        status: { type: "complete", reason: "unknown" },
-        ...result,
-      });
+
+      if (message.status.type === "running") {
+        updateMessage({
+          status: { type: "complete", reason: "unknown" },
+        });
+      } else {
+        // notify subscribers that isRunning is now false
+        this.notifySubscribers();
+      }
     } catch (e) {
       this.abortController = null;