@xsai/stream-text 0.0.27 → 0.0.29

package/dist/index.d.ts

@@ -1,6 +1,23 @@
-import { ChatOptions, FinishReason } from '@xsai/shared-chat';
+import { FinishReason, Usage, ChatOptions } from '@xsai/shared-chat';
 
+interface ChunkResult {
+    choices: {
+        delta: {
+            content: string;
+            role: 'assistant';
+        };
+        finish_reason?: FinishReason;
+        index: number;
+    }[];
+    created: number;
+    id: string;
+    model: string;
+    object: 'chat.completion.chunk';
+    system_fingerprint: string;
+    usage?: Usage;
+}
 interface StreamTextOptions extends ChatOptions {
+    onChunk?: (chunk: ChunkResult) => Promise<void> | void;
     /** if you want to disable stream, use `@xsai/generate-{text,object}` */
     stream?: never;
     streamOptions?: {
@@ -12,36 +29,15 @@ interface StreamTextOptions extends ChatOptions {
         usage?: boolean;
     };
 }
-interface StreamTextResponseUsage {
-    completion_tokens: number;
-    prompt_tokens: number;
-    total_tokens: number;
-}
 interface StreamTextResult {
-    chunkStream: ReadableStream<StreamTextResponse>;
+    chunkStream: ReadableStream<ChunkResult>;
     finishReason?: FinishReason;
     textStream: ReadableStream<string>;
-    usage?: StreamTextResponseUsage;
-}
-interface StreamTextResponse {
-    choices: {
-        delta: {
-            content: string;
-            role: 'assistant';
-        };
-        finish_reason?: FinishReason;
-        index: number;
-    }[];
-    created: number;
-    id: string;
-    model: string;
-    object: 'chat.completion.chunk';
-    system_fingerprint: string;
-    usage?: StreamTextResponseUsage;
+    usage?: Usage;
 }
 /**
  * @experimental WIP, does not support function calling (tools).
  */
 declare const streamText: (options: StreamTextOptions) => Promise<StreamTextResult>;
 
-export { type StreamTextOptions, type StreamTextResponse, type StreamTextResponseUsage, type StreamTextResult, streamText as default, streamText };
+export { type ChunkResult, type StreamTextOptions, type StreamTextResult, streamText as default, streamText };

package/dist/index.js

@@ -1,7 +1,6 @@
 import { chat } from '@xsai/shared-chat';
 
-const dataHeaderPrefix = "data: ";
-const dataErrorPrefix = `{"error":`;
+const chunkHeaderPrefix = "data:";
 const streamText = async (options) => await chat({
   ...options,
   stream: true
@@ -9,33 +8,40 @@ const streamText = async (options) => await chat({
   const decoder = new TextDecoder();
   let finishReason;
   let usage;
+  const processLine = async (line, controller) => {
+    if (!line || !line.startsWith(chunkHeaderPrefix))
+      return;
+    const content = line.slice(chunkHeaderPrefix.length);
+    const data = content.startsWith(" ") ? content.slice(1) : content;
+    if (data === "[DONE]") {
+      controller.terminate();
+      return true;
+    }
+    if (data.startsWith("{") && data.includes('"error":')) {
+      controller.error(new Error(`Error from server: ${data}`));
+      return true;
+    }
+    const chunk = JSON.parse(data);
+    controller.enqueue(chunk);
+    if (options.onChunk)
+      await options.onChunk(chunk);
+    if (chunk.choices[0].finish_reason) {
+      finishReason = chunk.choices[0].finish_reason;
+    }
+    if (chunk.usage) {
+      usage = chunk.usage;
+    }
+  };
   let buffer = "";
   const rawChunkStream = res.body.pipeThrough(new TransformStream({
-    transform: (chunk, controller) => {
-      buffer += decoder.decode(chunk);
-      const lines = buffer.split("\n\n");
+    transform: async (chunk, controller) => {
+      const text = decoder.decode(chunk, { stream: true });
+      buffer += text;
+      const lines = buffer.split("\n");
       buffer = lines.pop() || "";
       for (const line of lines) {
-        if (!line || !line.startsWith(dataHeaderPrefix)) {
-          continue;
-        }
-        if (line.startsWith(dataErrorPrefix)) {
-          controller.error(new Error(`Error from server: ${line}`));
-          break;
-        }
-        const lineWithoutPrefix = line.slice(dataHeaderPrefix.length);
-        if (lineWithoutPrefix === "[DONE]") {
-          controller.terminate();
+        if (await processLine(line, controller))
           break;
-        }
-        const data = JSON.parse(lineWithoutPrefix);
-        controller.enqueue(data);
-        if (data.choices[0].finish_reason) {
-          finishReason = data.choices[0].finish_reason;
-        }
-        if (data.usage) {
-          usage = data.usage;
-        }
       }
     }
   }));

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@xsai/stream-text",
-  "version": "0.0.27",
+  "version": "0.0.29",
   "type": "module",
   "author": "Moeru AI",
   "license": "MIT",