From 4d5a1e4a4e9dc00b636e422456039ecb989b5de5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=97=8D+85CD?= <50108258+kwaa@users.noreply.github.com>
Date: Wed, 15 Jan 2025 16:32:14 +0800
Subject: [PATCH] feat(stream-text): add onChunk() (#27)

* feat(stream-text): add onChunk()

* chore(stream-text): update test

* chore(stream-text): update test snapshot
---
 packages/stream-text/src/index.ts             | 28 +++++++++++--------
 .../test/__snapshots__/index.test.ts.snap     |  2 ++
 packages/stream-text/test/index.test.ts       |  8 +++++-
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/packages/stream-text/src/index.ts b/packages/stream-text/src/index.ts
index 38af4d7..794317c 100644
--- a/packages/stream-text/src/index.ts
+++ b/packages/stream-text/src/index.ts
@@ -5,6 +5,7 @@ import {
 } from '@xsai/shared-chat'
 
 export interface StreamTextOptions extends ChatOptions {
+  onChunk?: (chunk: StreamTextResponse) => Promise<void> | void
   /** if you want to disable stream, use `@xsai/generate-{text,object}` */
   stream?: never
   streamOptions?: {
@@ -48,8 +49,8 @@ export interface StreamTextResponse {
   usage?: StreamTextResponseUsage
 }
 
-const dataHeaderPrefix = 'data: '
-const dataErrorPrefix = `{"error":`
+const chunkHeaderPrefix = 'data: '
+const chunkErrorPrefix = `{"error":`
 
 /**
  * @experimental WIP, does not support function calling (tools).
@@ -66,7 +67,7 @@ export const streamText = async (options: StreamTextOptions): Promise<StreamText
 
   // null body handled by import('@xsai/shared-chat').chat()
   const rawChunkStream = res.body!.pipeThrough(new TransformStream({
-    transform: (chunk, controller) => {
+    transform: async (chunk, controller) => {
       buffer += decoder.decode(chunk)
       const lines = buffer.split('\n\n')
       buffer = lines.pop() || ''
@@ -75,30 +76,33 @@ export const streamText = async (options: StreamTextOptions): Promise<StreamText
         // Some cases:
         // - Empty chunk
         // - :ROUTER PROCESSING from OpenRouter
-        if (!line || !line.startsWith(dataHeaderPrefix)) {
+        if (!line || !line.startsWith(chunkHeaderPrefix)) {
           continue
         }
 
-        if (line.startsWith(dataErrorPrefix)) {
+        if (line.startsWith(chunkErrorPrefix)) {
           // About controller error: https://developer.mozilla.org/en-US/docs/Web/API/TransformStreamDefaultController/error
           controller.error(new Error(`Error from server: ${line}`))
           break
         }
 
-        const lineWithoutPrefix = line.slice(dataHeaderPrefix.length)
+        const lineWithoutPrefix = line.slice(chunkHeaderPrefix.length)
         if (lineWithoutPrefix === '[DONE]') {
           controller.terminate()
           break
         }
 
-        const data: StreamTextResponse = JSON.parse(lineWithoutPrefix)
-        controller.enqueue(data)
+        const chunk: StreamTextResponse = JSON.parse(lineWithoutPrefix)
+        controller.enqueue(chunk)
 
-        if (data.choices[0].finish_reason) {
-          finishReason = data.choices[0].finish_reason
+        if (options.onChunk)
+          await options.onChunk(chunk)
+
+        if (chunk.choices[0].finish_reason) {
+          finishReason = chunk.choices[0].finish_reason
         }
-        if (data.usage) {
-          usage = data.usage
+        if (chunk.usage) {
+          usage = chunk.usage
         }
       }
     },
diff --git a/packages/stream-text/test/__snapshots__/index.test.ts.snap b/packages/stream-text/test/__snapshots__/index.test.ts.snap
index d29a26f..68076ea 100644
--- a/packages/stream-text/test/__snapshots__/index.test.ts.snap
+++ b/packages/stream-text/test/__snapshots__/index.test.ts.snap
@@ -185,3 +185,5 @@ exports[`@xsai/stream-text > the-quick-brown-fox 2`] = `
   },
 ]
 `;
+
+exports[`@xsai/stream-text > the-quick-brown-fox 3`] = `11`;
diff --git a/packages/stream-text/test/index.test.ts b/packages/stream-text/test/index.test.ts
index a66eb75..b9cd98a 100644
--- a/packages/stream-text/test/index.test.ts
+++ b/packages/stream-text/test/index.test.ts
@@ -31,6 +31,9 @@ describe('@xsai/stream-text', () => {
   })
 
   it('the-quick-brown-fox', async () => {
+    let onChunkCount = 0
+    let chunkCount = 0
+
     const { chunkStream, textStream } = await streamText({
       ...ollama.chat('llama3.2'),
       messages: [
@@ -43,6 +46,7 @@ describe('@xsai/stream-text', () => {
           role: 'user',
         },
       ],
+      onChunk: () => { onChunkCount++ },
     })
 
     const chunk: StreamTextResponse[] = []
@@ -54,6 +58,7 @@ describe('@xsai/stream-text', () => {
         created: undefined,
         id: undefined,
       }))
+      chunkCount++
     }
 
     for await (const textPart of textStream) {
@@ -62,8 +67,9 @@ describe('@xsai/stream-text', () => {
 
     expect(text.join('')).toBe('The quick brown fox jumps over the lazy dog.')
     expect(text).toMatchSnapshot()
-
     expect(chunk).toMatchSnapshot()
+
+    expect(onChunkCount).toMatchSnapshot(chunkCount)
   })
 
   // TODO: error handling