Add types for data stream lines. (#741)

Co-authored-by: Max Leiter <max.leiter@vercel.com>

Author: lgrammel

.changeset/great-boxes-glow.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+experimental_StreamData/StreamingReactResponse: optimize parsing, improve types

examples/next-openai/app/stream-react-response/action.tsx

@@ -109,7 +109,7 @@ export async function handler({ messages }: { messages: Message[] }) {
     data,
     ui({ content, data }) {
       if (data != null) {
-        const value = (data as JSONValue[])[0] as any;
+        const value = data[0] as any;
 
         switch (value.type) {
           case 'weather': {

packages/core/react/parse-complex-response.test.ts

@@ -29,6 +29,8 @@ describe('parseComplexResponse function', () => {
       reader: createTestReader(['0:"Hello"\n']),
       abortControllerRef: { current: new AbortController() },
       update: mockUpdate,
+      generateId: () => 'test-id',
+      getCurrentDate: () => new Date(0),
     });
 
     const expectedMessage = assistantTextMessage('Hello');
@@ -37,9 +39,18 @@ describe('parseComplexResponse function', () => {
     expect(mockUpdate).toHaveBeenCalledTimes(1);
     expect(mockUpdate.mock.calls[0][0]).toEqual([expectedMessage]);
 
-    // check the prefix map:
-    expect(result).toHaveProperty('text');
-    expect(result.text).toEqual(expectedMessage);
+    // check the result
+    expect(result).toEqual({
+      messages: [
+        {
+          content: 'Hello',
+          createdAt: new Date(0),
+          id: 'test-id',
+          role: 'assistant',
+        },
+      ],
+      data: [],
+    });
   });
 
   it('should parse a sequence of text messages', async () => {
@@ -54,6 +65,8 @@ describe('parseComplexResponse function', () => {
       ]),
       abortControllerRef: { current: new AbortController() },
       update: mockUpdate,
+      generateId: () => 'test-id',
+      getCurrentDate: () => new Date(0),
     });
 
     // check the mockUpdate call:
@@ -71,20 +84,31 @@ describe('parseComplexResponse function', () => {
       assistantTextMessage('Hello, world.'),
     ]);
 
-    // check the prefix map:
-    expect(result).toHaveProperty('text');
-    expect(result.text).toEqual(assistantTextMessage('Hello, world.'));
+    // check the result
+    expect(result).toEqual({
+      messages: [
+        {
+          content: 'Hello, world.',
+          createdAt: new Date(0),
+          id: 'test-id',
+          role: 'assistant',
+        },
+      ],
+      data: [],
+    });
   });
 
   it('should parse a function call', async () => {
     const mockUpdate = jest.fn();
 
     const result = await parseComplexResponse({
       reader: createTestReader([
-        '1:"{\\"function_call\\": {\\"name\\": \\"get_current_weather\\", \\"arguments\\": \\"{\\\\n\\\\\\"location\\\\\\": \\\\\\"Charlottesville, Virginia\\\\\\",\\\\n\\\\\\"format\\\\\\": \\\\\\"celsius\\\\\\"\\\\n}\\"}}"\n',
+        '1:{"function_call":{"name":"get_current_weather","arguments":"{\\n\\"location\\": \\"Charlottesville, Virginia\\",\\n\\"format\\": \\"celsius\\"\\n}"}}\n',
       ]),
       abortControllerRef: { current: new AbortController() },
       update: mockUpdate,
+      generateId: () => 'test-id',
+      getCurrentDate: () => new Date(0),
     });
 
     // check the mockUpdate call:
@@ -104,21 +128,24 @@ describe('parseComplexResponse function', () => {
       },
     ]);
 
-    // check the prefix map:
-    expect(result.function_call).toEqual({
-      id: expect.any(String),
-      role: 'assistant',
-      content: '',
-      name: 'get_current_weather',
-      function_call: {
-        name: 'get_current_weather',
-        arguments:
-          '{\n"location": "Charlottesville, Virginia",\n"format": "celsius"\n}',
-      },
-      createdAt: expect.any(Date),
+    // check the result
+    expect(result).toEqual({
+      messages: [
+        {
+          content: '',
+          createdAt: new Date(0),
+          id: 'test-id',
+          role: 'assistant',
+          function_call: {
+            name: 'get_current_weather',
+            arguments:
+              '{\n"location": "Charlottesville, Virginia",\n"format": "celsius"\n}',
+          },
+          name: 'get_current_weather',
+        },
+      ],
+      data: [],
     });
-    expect(result).not.toHaveProperty('text');
-    expect(result).not.toHaveProperty('data');
   });
 
   it('should parse a combination of a data and a text message', async () => {
@@ -127,31 +154,74 @@ describe('parseComplexResponse function', () => {
     // Execute the parser function
     const result = await parseComplexResponse({
       reader: createTestReader([
-        '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+        '2:[{"t1":"v1"}]\n',
         '0:"Sample text message."\n',
       ]),
       abortControllerRef: { current: new AbortController() },
       update: mockUpdate,
+      generateId: () => 'test-id',
+      getCurrentDate: () => new Date(0),
     });
 
-    const expectedData = [{ t1: 'v1' }];
-
     // check the mockUpdate call:
     expect(mockUpdate).toHaveBeenCalledTimes(2);
 
     expect(mockUpdate.mock.calls[0][0]).toEqual([]);
-    expect(mockUpdate.mock.calls[0][1]).toEqual(expectedData);
+    expect(mockUpdate.mock.calls[0][1]).toEqual([{ t1: 'v1' }]);
 
     expect(mockUpdate.mock.calls[1][0]).toEqual([
       assistantTextMessage('Sample text message.'),
     ]);
-    expect(mockUpdate.mock.calls[1][1]).toEqual(expectedData);
+    expect(mockUpdate.mock.calls[1][1]).toEqual([{ t1: 'v1' }]);
+
+    // check the result
+    expect(result).toEqual({
+      messages: [
+        {
+          content: 'Sample text message.',
+          createdAt: new Date(0),
+          id: 'test-id',
+          role: 'assistant',
+        },
+      ],
+      data: [{ t1: 'v1' }],
+    });
+  });
 
-    // check the prefix map:
-    expect(result).toHaveProperty('data');
-    expect(result.data).toEqual(expectedData);
+  it('should parse multiple data messages incl. primitive values', async () => {
+    const mockUpdate = jest.fn();
 
-    expect(result).toHaveProperty('text');
-    expect(result.text).toEqual(assistantTextMessage('Sample text message.'));
+    // Execute the parser function
+    const result = await parseComplexResponse({
+      reader: createTestReader([
+        '2:[{"t1":"v1"}, 3]\n',
+        '2:[null,false,"text"]\n',
+      ]),
+      abortControllerRef: { current: new AbortController() },
+      update: mockUpdate,
+      generateId: () => 'test-id',
+      getCurrentDate: () => new Date(0),
+    });
+
+    // check the mockUpdate call:
+    expect(mockUpdate).toHaveBeenCalledTimes(2);
+
+    expect(mockUpdate.mock.calls[0][0]).toEqual([]);
+    expect(mockUpdate.mock.calls[0][1]).toEqual([{ t1: 'v1' }, 3]);
+
+    expect(mockUpdate.mock.calls[1][0]).toEqual([]);
+    expect(mockUpdate.mock.calls[1][1]).toEqual([
+      { t1: 'v1' },
+      3,
+      null,
+      false,
+      'text',
+    ]);
+
+    // check the result
+    expect(result).toEqual({
+      messages: [],
+      data: [{ t1: 'v1' }, 3, null, false, 'text'],
+    });
   });
 });

packages/core/react/parse-complex-response.ts

@@ -1,32 +1,40 @@
 import type { FunctionCall, JSONValue, Message } from '../shared/types';
-import { nanoid, createChunkDecoder } from '../shared/utils';
+import { createChunkDecoder, nanoid } from '../shared/utils';
 
 type PrefixMap = {
   text?: Message;
-  function_call?:
-    | string
-    | Pick<Message, 'function_call' | 'role' | 'content' | 'name'>;
-  data?: JSONValue[];
+  function_call?: Message & {
+    role: 'assistant';
+    function_call: FunctionCall;
+  };
+  data: JSONValue[];
 };
 
 export async function parseComplexResponse({
   reader,
   abortControllerRef,
   update,
   onFinish,
+  generateId = nanoid,
+  getCurrentDate = () => new Date(),
 }: {
   reader: ReadableStreamDefaultReader<Uint8Array>;
   abortControllerRef?: {
     current: AbortController | null;
   };
   update: (merged: Message[], data: JSONValue[] | undefined) => void;
   onFinish?: (prefixMap: PrefixMap) => void;
+  generateId?: () => string;
+  getCurrentDate?: () => Date;
 }) {
+  const createdAt = getCurrentDate();
+
   const decode = createChunkDecoder(true);
-  const createdAt = new Date();
-  const prefixMap: PrefixMap = {};
+  const prefixMap: PrefixMap = {
+    data: [],
+  };
   const NEWLINE = '\n'.charCodeAt(0);
-  let chunks: Uint8Array[] = [];
+  const chunks: Uint8Array[] = [];
   let totalLength = 0;
 
   while (true) {
@@ -73,7 +81,7 @@ export async function parseComplexResponse({
           };
         } else {
           prefixMap['text'] = {
-            id: nanoid(),
+            id: generateId(),
             role: 'assistant',
             content: value,
             createdAt,
@@ -84,46 +92,30 @@ export async function parseComplexResponse({
       let functionCallMessage: Message | null = null;
 
       if (type === 'function_call') {
-        prefixMap['function_call'] = value;
-
-        let functionCall = prefixMap['function_call'];
-        // Ensure it hasn't been parsed
-        if (functionCall && typeof functionCall === 'string') {
-          const parsedFunctionCall: FunctionCall = JSON.parse(
-            functionCall as string,
-          ).function_call;
-
-          functionCallMessage = {
-            id: nanoid(),
-            role: 'assistant',
-            content: '',
-            function_call: parsedFunctionCall,
-            name: parsedFunctionCall.name,
-            createdAt,
-          };
-
-          prefixMap['function_call'] = functionCallMessage as any;
-        }
+        prefixMap['function_call'] = {
+          id: generateId(),
+          role: 'assistant',
+          content: '',
+          function_call: value.function_call,
+          name: value.function_call.name,
+          createdAt,
+        };
+
+        functionCallMessage = prefixMap['function_call'];
       }
 
       if (type === 'data') {
-        const parsedValue = JSON.parse(value);
-        if (prefixMap['data']) {
-          prefixMap['data'] = [...prefixMap['data'], ...parsedValue];
-        } else {
-          prefixMap['data'] = parsedValue;
-        }
+        prefixMap['data'].push(...value);
       }
 
-      const data = prefixMap['data'];
       const responseMessage = prefixMap['text'];
 
       // We add function calls and response messages to the messages[], but data is its own thing
       const merged = [functionCallMessage, responseMessage].filter(
         Boolean,
       ) as Message[];
 
-      update(merged, data);
+      update(merged, [...prefixMap['data']]); // make a copy of the data array
 
       // The request has been aborted, stop reading the stream.
       // If abortControllerRef is undefined, this is intentionally not executed.
@@ -136,5 +128,10 @@ export async function parseComplexResponse({
 
   onFinish?.(prefixMap);
 
-  return prefixMap;
+  return {
+    messages: [prefixMap.text, prefixMap.function_call].filter(
+      Boolean,
+    ) as Message[],
+    data: prefixMap.data,
+  };
 }

packages/core/react/use-chat.ts

@@ -199,33 +199,22 @@ const getStreamedResponse = async (
   }
 
   const isComplexMode = res.headers.get(COMPLEX_HEADER) === 'true';
-  let responseMessages: Message[] = [];
   const reader = res.body.getReader();
 
-  // END TODO-STREAMDATA
-  let responseData: any = [];
-
   if (isComplexMode) {
-    const prefixMap = await parseComplexResponse({
+    return await parseComplexResponse({
       reader,
       abortControllerRef,
       update(merged, data) {
         mutate([...chatRequest.messages, ...merged], false);
         mutateStreamData([...(existingData || []), ...(data || [])], false);
       },
+      onFinish(prefixMap) {
+        if (onFinish && prefixMap.text != null) {
+          onFinish(prefixMap.text);
+        }
+      },
     });
-
-    for (const [type, item] of Object.entries(prefixMap)) {
-      if (onFinish && type === 'text') {
-        onFinish(item as Message);
-      }
-      if (type === 'data') {
-        responseData.push(item);
-      } else {
-        responseMessages.push(item as Message);
-      }
-    }
-    return { messages: responseMessages, data: responseData };
   } else {
     const createdAt = new Date();
     const decode = createChunkDecoder(false);

packages/core/shared/stream-parts.test.ts

@@ -0,0 +1,65 @@
+import { formatStreamPart, parseStreamPart } from './stream-parts';
+
+describe('stream-parts', () => {
+  describe('formatStreamPart', () => {
+    it('should escape newlines in text', () => {
+      expect(formatStreamPart('text', 'value\nvalue')).toEqual(
+        '0:"value\\nvalue"\n',
+      );
+    });
+
+    it('should escape newlines in data objects', () => {
+      expect(formatStreamPart('data', [{ test: 'value\nvalue' }])).toEqual(
+        '2:[{"test":"value\\nvalue"}]\n',
+      );
+    });
+  });
+
+  describe('parseStreamPart', () => {
+    it('should parse a text line', () => {
+      const input = '0:"Hello, world!"';
+
+      expect(parseStreamPart(input)).toEqual({
+        type: 'text',
+        value: 'Hello, world!',
+      });
+    });
+
+    it('should parse a function call line', () => {
+      const input =
+        '1:{"function_call": {"name":"get_current_weather","arguments":"{\\"location\\": \\"Charlottesville, Virginia\\",\\"format\\": \\"celsius\\"}"}}';
+
+      expect(parseStreamPart(input)).toEqual({
+        type: 'function_call',
+        value: {
+          function_call: {
+            name: 'get_current_weather',
+            arguments:
+              '{"location": "Charlottesville, Virginia","format": "celsius"}',
+          },
+        },
+      });
+    });
+
+    it('should parse a data line', () => {
+      const input = '2:[{"test":"value"}]';
+      const expectedOutput = { type: 'data', value: [{ test: 'value' }] };
+      expect(parseStreamPart(input)).toEqual(expectedOutput);
+    });
+
+    it('should throw an error if the input does not contain a colon separator', () => {
+      const input = 'invalid stream string';
+      expect(() => parseStreamPart(input)).toThrow();
+    });
+
+    it('should throw an error if the input contains an invalid type', () => {
+      const input = '55:test';
+      expect(() => parseStreamPart(input)).toThrow();
+    });
+
+    it("should throw error if the input's JSON is invalid", () => {
+      const input = '0:{"test":"value"';
+      expect(() => parseStreamPart(input)).toThrow();
+    });
+  });
+});

packages/core/shared/stream-parts.ts

@@ -0,0 +1,148 @@
+import { FunctionCall, JSONValue } from './types';
+import { StreamString } from './utils';
+
+export interface StreamPart<CODE extends string, NAME extends string, TYPE> {
+  code: CODE;
+  name: NAME;
+  parse: (value: JSONValue) => { type: NAME; value: TYPE };
+}
+
+export const textStreamPart: StreamPart<'0', 'text', string> = {
+  code: '0',
+  name: 'text',
+  parse: (value: JSONValue) => {
+    if (typeof value !== 'string') {
+      throw new Error('"text" parts expect a string value.');
+    }
+    return { type: 'text', value };
+  },
+};
+
+export const functionCallStreamPart: StreamPart<
+  '1',
+  'function_call',
+  { function_call: FunctionCall }
+> = {
+  code: '1',
+  name: 'function_call',
+  parse: (value: JSONValue) => {
+    if (
+      value == null ||
+      typeof value !== 'object' ||
+      !('function_call' in value)
+    ) {
+      throw new Error(
+        '"function_call" parts expect an object with a "function_call" property.',
+      );
+    }
+
+    const functionCall = value.function_call;
+
+    if (
+      functionCall == null ||
+      typeof functionCall !== 'object' ||
+      !('name' in functionCall) ||
+      !('arguments' in functionCall)
+    ) {
+      throw new Error(
+        '"function_call" parts expect an object with a "name" and "arguments" property.',
+      );
+    }
+
+    return {
+      type: 'function_call',
+      value: value as unknown as { function_call: FunctionCall },
+    };
+  },
+};
+
+export const dataStreamPart: StreamPart<'2', 'data', Array<JSONValue>> = {
+  code: '2',
+  name: 'data',
+  parse: (value: JSONValue) => {
+    if (!Array.isArray(value)) {
+      throw new Error('"data" parts expect an array value.');
+    }
+
+    return { type: 'data', value };
+  },
+};
+
+const streamParts = [
+  textStreamPart,
+  functionCallStreamPart,
+  dataStreamPart,
+] as const;
+
+// union type of all stream parts
+type StreamParts =
+  | typeof textStreamPart
+  | typeof functionCallStreamPart
+  | typeof dataStreamPart;
+
+/**
+ * Maps the type of a stream part to its value type.
+ */
+type StreamPartValueType = {
+  [P in StreamParts as P['name']]: ReturnType<P['parse']>['value'];
+};
+
+export type StreamPartType =
+  | ReturnType<typeof textStreamPart.parse>
+  | ReturnType<typeof functionCallStreamPart.parse>
+  | ReturnType<typeof dataStreamPart.parse>;
+
+export const streamPartsByCode = {
+  [textStreamPart.code]: textStreamPart,
+  [functionCallStreamPart.code]: functionCallStreamPart,
+  [dataStreamPart.code]: dataStreamPart,
+} as const;
+
+export const validCodes = streamParts.map(part => part.code);
+
+/**
+ * Parses a stream part from a string.
+ *
+ * @param line The string to parse.
+ * @returns The parsed stream part.
+ * @throws An error if the string cannot be parsed.
+ */
+export const parseStreamPart = (line: string): StreamPartType => {
+  const firstSeperatorIndex = line.indexOf(':');
+
+  if (firstSeperatorIndex === -1) {
+    throw new Error('Failed to parse stream string. No seperator found.');
+  }
+
+  const prefix = line.slice(0, firstSeperatorIndex);
+
+  if (!validCodes.includes(prefix as keyof typeof streamPartsByCode)) {
+    throw new Error(`Failed to parse stream string. Invalid code ${prefix}.`);
+  }
+
+  const code = prefix as keyof typeof streamPartsByCode;
+
+  const textValue = line.slice(firstSeperatorIndex + 1);
+  const jsonValue: JSONValue = JSON.parse(textValue);
+
+  return streamPartsByCode[code].parse(jsonValue);
+};
+
+/**
+ * Prepends a string with a prefix from the `StreamChunkPrefixes`, JSON-ifies it,
+ * and appends a new line.
+ *
+ * It ensures type-safety for the part type and value.
+ */
+export function formatStreamPart<T extends keyof StreamPartValueType>(
+  type: T,
+  value: StreamPartValueType[T],
+): StreamString {
+  const streamPart = streamParts.find(part => part.name === type);
+
+  if (!streamPart) {
+    throw new Error(`Invalid stream part type: ${type}`);
+  }
+
+  return `${streamPart.code}:${JSON.stringify(value)}\n`;
+}

packages/core/shared/utils.test.ts

@@ -1,13 +1,13 @@
-import { createChunkDecoder, getStreamString } from './utils';
-import { getStreamStringTypeAndValue } from './utils';
+import { formatStreamPart } from './stream-parts';
+import { createChunkDecoder } from './utils';
 
 describe('utils', () => {
   describe('createChunkDecoder', () => {
     it('should correctly decode text chunk in complex mode', () => {
       const decoder = createChunkDecoder(true);
 
       const encoder = new TextEncoder();
-      const chunk = encoder.encode(getStreamString('text', 'Hello, world!'));
+      const chunk = encoder.encode(formatStreamPart('text', 'Hello, world!'));
       const values = decoder(chunk);
 
       expect(values).toStrictEqual([{ type: 'text', value: 'Hello, world!' }]);
@@ -24,22 +24,29 @@ describe('utils', () => {
 
       const encoder = new TextEncoder();
       const chunk = encoder.encode(
-        getStreamString('function_call', functionCall),
+        formatStreamPart('function_call', {
+          function_call: functionCall,
+        }),
       );
       const values = decoder(chunk);
 
       expect(values).toStrictEqual([
-        { type: 'function_call', value: functionCall },
+        {
+          type: 'function_call',
+          value: {
+            function_call: functionCall,
+          },
+        },
       ]);
     });
 
     it('should correctly decode data chunk in complex mode', () => {
-      const data = { test: 'value' };
+      const data = [{ test: 'value' }];
 
       const decoder = createChunkDecoder(true);
 
       const encoder = new TextEncoder();
-      const chunk = encoder.encode(getStreamString('data', data));
+      const chunk = encoder.encode(formatStreamPart('data', data));
       const values = decoder(chunk);
 
       expect(values).toStrictEqual([{ type: 'data', value: data }]);
@@ -57,10 +64,10 @@ describe('utils', () => {
 
       const enqueuedChunks = [];
       enqueuedChunks.push(
-        encoder.encode(getStreamString('text', normalDecode(chunk1))),
+        encoder.encode(formatStreamPart('text', normalDecode(chunk1))),
       );
       enqueuedChunks.push(
-        encoder.encode(getStreamString('text', normalDecode(chunk2))),
+        encoder.encode(formatStreamPart('text', normalDecode(chunk2))),
       );
 
       let fullDecodedString = '';
@@ -91,40 +98,4 @@ describe('utils', () => {
       expect(values + secondValues).toBe('♥');
     });
   });
-
-  describe('getStreamStringTypeAndValue', () => {
-    it('should correctly parse a text stream string', () => {
-      const input = '0:Hello, world!';
-
-      expect(getStreamStringTypeAndValue(input)).toEqual({
-        type: 'text',
-        value: 'Hello, world!',
-      });
-    });
-
-    it('should correctly parse a function call stream string', () => {
-      const input =
-        '1:{"name":"get_current_weather","arguments":"{\\"location\\": \\"Charlottesville, Virginia\\",\\"format\\": \\"celsius\\"}"}';
-
-      expect(getStreamStringTypeAndValue(input)).toEqual({
-        type: 'function_call',
-        value: {
-          name: 'get_current_weather',
-          arguments:
-            '{"location": "Charlottesville, Virginia","format": "celsius"}',
-        },
-      });
-    });
-
-    it('should correctly parse a data stream string', () => {
-      const input = '2:{"test":"value"}';
-      const expectedOutput = { type: 'data', value: { test: 'value' } };
-      expect(getStreamStringTypeAndValue(input)).toEqual(expectedOutput);
-    });
-
-    it('should throw an error if the input is not a valid stream string', () => {
-      const input = 'invalid stream string';
-      expect(() => getStreamStringTypeAndValue(input)).toThrow();
-    });
-  });
 });

packages/core/shared/utils.ts

@@ -1,5 +1,13 @@
 import { customAlphabet } from 'nanoid/non-secure';
 import { JSONValue } from './types';
+import {
+  StreamPartType,
+  dataStreamPart,
+  functionCallStreamPart,
+  parseStreamPart,
+  streamPartsByCode,
+  textStreamPart,
+} from './stream-parts';
 
 // 7-character random string
 export const nanoid = customAlphabet(
@@ -13,19 +21,13 @@ function createChunkDecoder(
   complex: false,
 ): (chunk: Uint8Array | undefined) => string;
 // complex decoder signature:
-function createChunkDecoder(complex: true): (chunk: Uint8Array | undefined) => {
-  type: keyof typeof StreamStringPrefixes;
-  value: string;
-}[];
+function createChunkDecoder(
+  complex: true,
+): (chunk: Uint8Array | undefined) => StreamPartType[];
 // combined signature for when the client calls this function with a boolean:
-function createChunkDecoder(complex?: boolean): (
-  chunk: Uint8Array | undefined,
-) =>
-  | {
-      type: keyof typeof StreamStringPrefixes;
-      value: string;
-    }[]
-  | string;
+function createChunkDecoder(
+  complex?: boolean,
+): (chunk: Uint8Array | undefined) => StreamPartType[] | string;
 function createChunkDecoder(complex?: boolean) {
   const decoder = new TextDecoder();
 
@@ -42,7 +44,7 @@ function createChunkDecoder(complex?: boolean) {
       .split('\n')
       .filter(line => line !== ''); // splitting leaves an empty string at the end
 
-    return decoded.map(getStreamStringTypeAndValue).filter(Boolean) as any;
+    return decoded.map(parseStreamPart).filter(Boolean);
   };
 }
 
@@ -69,10 +71,9 @@ export { createChunkDecoder };
  *```
  */
 export const StreamStringPrefixes = {
-  text: 0,
-  function_call: 1,
-  data: 2,
-  // user_err: 3?
+  [textStreamPart.name]: textStreamPart.code,
+  [functionCallStreamPart.name]: functionCallStreamPart.code,
+  [dataStreamPart.name]: dataStreamPart.code,
 } as const;
 
 export const isStreamStringEqualToType = (
@@ -81,50 +82,9 @@ export const isStreamStringEqualToType = (
 ): value is StreamString =>
   value.startsWith(`${StreamStringPrefixes[type]}:`) && value.endsWith('\n');
 
-/**
- * Prepends a string with a prefix from the `StreamChunkPrefixes`, JSON-ifies it, and appends a new line.
- */
-export const getStreamString = (
-  type: keyof typeof StreamStringPrefixes,
-  value: JSONValue,
-): StreamString => `${StreamStringPrefixes[type]}:${JSON.stringify(value)}\n`;
-
 export type StreamString =
   `${(typeof StreamStringPrefixes)[keyof typeof StreamStringPrefixes]}:${string}\n`;
 
-export const getStreamStringTypeAndValue = (
-  line: string,
-): { type: keyof typeof StreamStringPrefixes; value: JSONValue } => {
-  const firstSeperatorIndex = line.indexOf(':');
-
-  if (firstSeperatorIndex === -1) {
-    throw new Error('Failed to parse stream string');
-  }
-
-  const prefix = line.slice(0, firstSeperatorIndex);
-  const type = Object.keys(StreamStringPrefixes).find(
-    key =>
-      StreamStringPrefixes[key as keyof typeof StreamStringPrefixes] ===
-      Number(prefix),
-  ) as keyof typeof StreamStringPrefixes;
-
-  const val = line.slice(firstSeperatorIndex + 1);
-
-  let parsedVal = val;
-
-  if (!val) {
-    return { type, value: '' };
-  }
-
-  try {
-    parsedVal = JSON.parse(val);
-  } catch (e) {
-    console.error('Failed to parse JSON value:', val);
-  }
-
-  return { type, value: parsedVal };
-};
-
 /**
  * A header sent to the client so it knows how to handle parsing the stream (as a deprecated text response or using the new prefixed protocol)
  */

packages/core/streams/cohere-stream.test.ts

@@ -91,7 +91,7 @@ describe('CohereStream', () => {
       const response = new StreamingTextResponse(stream, {}, data);
 
       expect(await readAllChunks(response)).toEqual([
-        '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+        '2:[{"t1":"v1"}]\n',
         '0:" Hello"\n',
         '0:","\n',
         '0:" world"\n',

packages/core/streams/huggingface-stream.test.ts

@@ -113,7 +113,7 @@ describe('HuggingFace stream', () => {
       const response = new StreamingTextResponse(stream, {}, data);
 
       expect(await readAllChunks(response)).toEqual([
-        '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+        '2:[{"t1":"v1"}]\n',
         '0:"Hello"\n',
         '0:","\n',
         '0:" world"\n',

packages/core/streams/langchain-stream.test.ts

@@ -156,7 +156,7 @@ describe('LangchainStream', () => {
         );
 
         expect(await readAllChunks(response)).toEqual([
-          '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+          '2:[{"t1":"v1"}]\n',
           '0:""\n',
           '0:"Hello"\n',
           '0:","\n',
@@ -270,7 +270,7 @@ describe('LangchainStream', () => {
         const response = new StreamingTextResponse(stream, {}, data);
 
         expect(await readAllChunks(response)).toEqual([
-          '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+          '2:[{"t1":"v1"}]\n',
           '0:""\n',
           '0:"Hello"\n',
           '0:","\n',

packages/core/streams/openai-stream.test.tsx

@@ -230,7 +230,7 @@ describe('OpenAIStream', () => {
       const chunks = await client.readAll();
 
       expect(chunks).toEqual([
-        '1:"{\\"function_call\\": {\\"name\\": \\"get_current_weather\\", \\"arguments\\": \\"{\\\\n\\\\\\"location\\\\\\": \\\\\\"Charlottesville, Virginia\\\\\\",\\\\n\\\\\\"format\\\\\\": \\\\\\"celsius\\\\\\"\\\\n}\\"}}"\n',
+        '1:{"function_call":{"name":"get_current_weather","arguments":"{\\n\\"location\\": \\"Charlottesville, Virginia\\",\\n\\"format\\": \\"celsius\\"\\n}"}}\n',
       ]);
     });
 
@@ -263,8 +263,8 @@ describe('OpenAIStream', () => {
       const chunks = await client.readAll();
 
       expect(chunks).toEqual([
-        '2:"[{\\"fn\\":\\"get_current_weather\\"}]"\n',
-        '1:"{\\"function_call\\": {\\"name\\": \\"get_current_weather\\", \\"arguments\\": \\"{\\\\n\\\\\\"location\\\\\\": \\\\\\"Charlottesville, Virginia\\\\\\",\\\\n\\\\\\"format\\\\\\": \\\\\\"celsius\\\\\\"\\\\n}\\"}}"\n',
+        '2:[{"fn":"get_current_weather"}]\n',
+        '1:{"function_call":{"name":"get_current_weather","arguments":"{\\n\\"location\\": \\"Charlottesville, Virginia\\",\\n\\"format\\": \\"celsius\\"\\n}"}}\n',
       ]);
     });
 
@@ -327,7 +327,7 @@ describe('OpenAIStream', () => {
       const chunks = await client.readAll();
 
       expect(chunks).toEqual([
-        '2:"[{\\"fn\\":\\"get_current_weather\\"}]"\n',
+        '2:[{"fn":"get_current_weather"}]\n',
         '0:"experimental_onFunctionCall-return-value"\n',
       ]);
     });
@@ -358,7 +358,7 @@ describe('OpenAIStream', () => {
       const chunks = await client.readAll();
 
       expect(chunks).toEqual([
-        '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+        '2:[{"t1":"v1"}]\n',
         '0:"Hello"\n',
         '0:","\n',
         '0:" world"\n',

packages/core/streams/openai-stream.ts

@@ -1,10 +1,11 @@
+import { formatStreamPart } from '../shared/stream-parts';
 import {
   CreateMessage,
   FunctionCall,
   JSONValue,
   Message,
 } from '../shared/types';
-import { createChunkDecoder, getStreamString } from '../shared/utils';
+import { createChunkDecoder } from '../shared/utils';
 
 import {
   AIStream,
@@ -487,7 +488,7 @@ function createFunctionCallTransformer(
       if (!isFunctionStreamingIn) {
         controller.enqueue(
           isComplexMode
-            ? textEncoder.encode(getStreamString('text', message))
+            ? textEncoder.encode(formatStreamPart('text', message))
             : chunk,
         );
         return;
@@ -546,7 +547,11 @@ function createFunctionCallTransformer(
             controller.enqueue(
               textEncoder.encode(
                 isComplexMode
-                  ? getStreamString('function_call', aggregatedResponse)
+                  ? formatStreamPart(
+                      'function_call',
+                      // parse to prevent double-encoding:
+                      JSON.parse(aggregatedResponse),
+                    )
                   : aggregatedResponse,
               ),
             );
@@ -555,7 +560,7 @@ function createFunctionCallTransformer(
             // The user returned a string, so we just return it as a message
             controller.enqueue(
               isComplexMode
-                ? textEncoder.encode(getStreamString('text', functionResponse))
+                ? textEncoder.encode(formatStreamPart('text', functionResponse))
                 : textEncoder.encode(functionResponse),
             );
             return;

packages/core/streams/replicate-stream.test.ts

@@ -104,7 +104,7 @@ describe('ReplicateStream', () => {
       const response = new StreamingTextResponse(stream, {}, data);
 
       expect(await readAllChunks(response)).toEqual([
-        '2:"[{\\"t1\\":\\"v1\\"}]"\n',
+        '2:[{"t1":"v1"}]\n',
         '0:" Hello,"\n',
         '0:" world"\n',
         '0:"."\n',

packages/core/streams/stream-data.ts

@@ -1,5 +1,5 @@
+import { formatStreamPart } from '../shared/stream-parts';
 import { JSONValue } from '../shared/types';
-import { getStreamString } from '../shared/utils';
 
 /**
  * A stream wrapper to send custom JSON-encoded data back to the client.
@@ -33,7 +33,7 @@ export class experimental_StreamData {
         // add buffered data to the stream
         if (self.data.length > 0) {
           const encodedData = self.encoder.encode(
-            getStreamString('data', JSON.stringify(self.data)),
+            formatStreamPart('data', self.data),
           );
           self.data = [];
           controller.enqueue(encodedData);
@@ -60,7 +60,7 @@ export class experimental_StreamData {
 
         if (self.data.length) {
           const encodedData = self.encoder.encode(
-            getStreamString('data', JSON.stringify(self.data)),
+            formatStreamPart('data', self.data),
           );
           controller.enqueue(encodedData);
         }
@@ -109,7 +109,7 @@ export function createStreamDataTransformer(
   return new TransformStream({
     transform: async (chunk, controller) => {
       const message = decoder.decode(chunk);
-      controller.enqueue(encoder.encode(getStreamString('text', message)));
+      controller.enqueue(encoder.encode(formatStreamPart('text', message)));
     },
   });
 }