Fix queue closing.

vercel · Jul 11, 2023 · 75aa96c · 75aa96c · vercel · Jul 11, 2023
1 parent a5d4a84
commit 75aa96c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/examples/ai-chat-next-js/src/pages/api/send-message.ts b/examples/ai-chat-next-js/src/pages/api/send-message.ts
@@ -1,4 +1,5 @@
 import {
+  CohereTextGenerationModel,
   LlamaCppTextGenerationModel,
   OpenAIChatMessage,
   OpenAIChatModel,
@@ -66,8 +67,12 @@ async function createTextStream(
   controller: AbortController
 ) {
   return streamText(
-    // You can user other models such as CohereTextGenerationModel here:
     new LlamaCppTextGenerationModel({}),
+    // You can also use other models such as CohereTextGenerationModel instead:
+    // new CohereTextGenerationModel({
+    //   model: "command-nightly",
+    //   maxTokens: 1000,
+    // }),
     [
       "You are an AI chat bot. " +
         "Follow the user's instructions carefully. Respond using markdown.",

diff --git a/src/model/text-streaming/TextDeltaEventSource.ts b/src/model/text-streaming/TextDeltaEventSource.ts
@@ -56,8 +56,13 @@ export function convertTextDeltaEventSourceToIterable(
         );
 
         queue.push(data.textDelta);
+
+        if (data.isFinished) {
+          queue.close();
+        }
       } catch (error) {
         options?.errorHandler(error);
+        queue.close();
       }
     },
   });