Fix trailing empty line after split in Decoder. (#684)

Co-authored-by: Max Leiter <max.leiter@vercel.com>
vercel · Oct 26, 2023 · 424d5ee · 424d5ee
1 parent 7f58d9b
commit 424d5ee
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 42 deletions.
diff --git a/.changeset/mean-beans-approve.md b/.changeset/mean-beans-approve.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+experimental_StreamData: fix trailing newline parsing bug in decoder
diff --git a/packages/core/react/parseComplexResponse.test.ts b/packages/core/react/parseComplexResponse.test.ts
@@ -34,9 +34,8 @@ describe('parseComplexResponse function', () => {
     const expectedMessage = assistantTextMessage('Hello');
 
     // check the mockUpdate call:
-    expect(mockUpdate).toHaveBeenCalledTimes(2);
+    expect(mockUpdate).toHaveBeenCalledTimes(1);
     expect(mockUpdate.mock.calls[0][0]).toEqual([expectedMessage]);
-    expect(mockUpdate.mock.calls[1][0]).toEqual([expectedMessage]);
 
     // check the prefix map:
     expect(result).toHaveProperty('text');
@@ -58,29 +57,17 @@ describe('parseComplexResponse function', () => {
     });
 
     // check the mockUpdate call:
-    expect(mockUpdate).toHaveBeenCalledTimes(8);
+    expect(mockUpdate).toHaveBeenCalledTimes(4);
     expect(mockUpdate.mock.calls[0][0]).toEqual([
       assistantTextMessage('Hello'),
     ]);
     expect(mockUpdate.mock.calls[1][0]).toEqual([
-      assistantTextMessage('Hello'),
-    ]);
-    expect(mockUpdate.mock.calls[2][0]).toEqual([
-      assistantTextMessage('Hello,'),
-    ]);
-    expect(mockUpdate.mock.calls[3][0]).toEqual([
       assistantTextMessage('Hello,'),
     ]);
-    expect(mockUpdate.mock.calls[4][0]).toEqual([
-      assistantTextMessage('Hello, world'),
-    ]);
-    expect(mockUpdate.mock.calls[5][0]).toEqual([
+    expect(mockUpdate.mock.calls[2][0]).toEqual([
       assistantTextMessage('Hello, world'),
     ]);
-    expect(mockUpdate.mock.calls[6][0]).toEqual([
-      assistantTextMessage('Hello, world.'),
-    ]);
-    expect(mockUpdate.mock.calls[7][0]).toEqual([
+    expect(mockUpdate.mock.calls[3][0]).toEqual([
       assistantTextMessage('Hello, world.'),
     ]);
 
@@ -101,7 +88,7 @@ describe('parseComplexResponse function', () => {
     });
 
     // check the mockUpdate call:
-    expect(mockUpdate).toHaveBeenCalledTimes(2);
+    expect(mockUpdate).toHaveBeenCalledTimes(1);
     expect(mockUpdate.mock.calls[0][0]).toEqual([
       {
         id: expect.any(String),
@@ -116,11 +103,8 @@ describe('parseComplexResponse function', () => {
         createdAt: expect.any(Date),
       },
     ]);
-    expect(mockUpdate.mock.calls[1][0]).toEqual([assistantTextMessage('')]);
 
     // check the prefix map:
-    expect(result).toHaveProperty('text');
-    expect(result.text).toEqual(assistantTextMessage(''));
     expect(result.function_call).toEqual({
       id: expect.any(String),
       role: 'assistant',
@@ -133,6 +117,7 @@ describe('parseComplexResponse function', () => {
       },
       createdAt: expect.any(Date),
     });
+    expect(result).not.toHaveProperty('text');
     expect(result).not.toHaveProperty('data');
   });
 
@@ -152,23 +137,15 @@ describe('parseComplexResponse function', () => {
     const expectedData = [{ t1: 'v1' }];
 
     // check the mockUpdate call:
-    expect(mockUpdate).toHaveBeenCalledTimes(4);
+    expect(mockUpdate).toHaveBeenCalledTimes(2);
 
     expect(mockUpdate.mock.calls[0][0]).toEqual([]);
     expect(mockUpdate.mock.calls[0][1]).toEqual(expectedData);
 
-    expect(mockUpdate.mock.calls[1][0]).toEqual([assistantTextMessage('')]);
-    expect(mockUpdate.mock.calls[1][1]).toEqual(expectedData);
-
-    expect(mockUpdate.mock.calls[2][0]).toEqual([
-      assistantTextMessage('Sample text message.'),
-    ]);
-    expect(mockUpdate.mock.calls[2][1]).toEqual(expectedData);
-
-    expect(mockUpdate.mock.calls[3][0]).toEqual([
+    expect(mockUpdate.mock.calls[1][0]).toEqual([
       assistantTextMessage('Sample text message.'),
     ]);
-    expect(mockUpdate.mock.calls[3][1]).toEqual(expectedData);
+    expect(mockUpdate.mock.calls[1][1]).toEqual(expectedData);
 
     // check the prefix map:
     expect(result).toHaveProperty('data');

diff --git a/packages/core/shared/utils.test.ts b/packages/core/shared/utils.test.ts
@@ -9,10 +9,7 @@ describe('utils', () => {
       const chunk = encoder.encode(getStreamString('text', 'Hello, world!'));
       const values = decoder(chunk);
 
-      expect(values).toStrictEqual([
-        { type: 'text', value: 'Hello, world!' },
-        { type: 'text', value: '' },
-      ]);
+      expect(values).toStrictEqual([{ type: 'text', value: 'Hello, world!' }]);
     });
 
     it('should correctly decode function chunk in complex mode', () => {
@@ -32,7 +29,6 @@ describe('utils', () => {
 
       expect(values).toStrictEqual([
         { type: 'function_call', value: functionCall },
-        { type: 'text', value: '' },
       ]);
     });
 
@@ -45,10 +41,7 @@ describe('utils', () => {
       const chunk = encoder.encode(getStreamString('data', data));
       const values = decoder(chunk);
 
-      expect(values).toStrictEqual([
-        { type: 'data', value: data },
-        { type: 'text', value: '' },
-      ]);
+      expect(values).toStrictEqual([{ type: 'data', value: data }]);
     });
 
     it('should correctly decode streamed utf8 chunks in complex mode', () => {

diff --git a/packages/core/shared/utils.ts b/packages/core/shared/utils.ts
@@ -37,7 +37,11 @@ function createChunkDecoder(complex?: boolean) {
   }
 
   return function (chunk: Uint8Array | undefined) {
-    const decoded = decoder.decode(chunk, { stream: true }).split('\n');
+    const decoded = decoder
+      .decode(chunk, { stream: true })
+      .split('\n')
+      .filter(line => line !== ''); // splitting leaves an empty string at the end
+
     return decoded.map(getStreamStringTypeAndValue).filter(Boolean) as any;
   };
 }