diff --git a/packages/stream-text/src/index.ts b/packages/stream-text/src/index.ts index 38af4d7e..794317c7 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 /** 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 { + 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 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 a66eb757..b9cd98a1 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