Expose StreamPart & utilities. (#1334)

lgrammel · web-flow · commit ea6b0e1569ff · 2024-04-12T15:57:38.000+02:00
diff --git a/.changeset/weak-bags-attack.md b/.changeset/weak-bags-attack.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+Expose formatStreamPart, parseStreamPart, and readDataStream helpers.
diff --git a/packages/core/shared/read-data-stream.ts b/packages/core/shared/read-data-stream.ts
@@ -16,6 +16,17 @@ function concatChunks(chunks: Uint8Array[], totalLength: number) {
   return concatenatedChunks;
 }
 
+/**
+Converts a ReadableStreamDefaultReader into an async generator that yields
+StreamPart objects.
+
+@param reader 
+       Reader for the stream to read from.
+@param isAborted
+       Optional function that returns true if the request has been aborted.
+       If the function returns true, the generator will stop reading the stream.
+       If the function is not provided, the generator will not stop reading the stream.
+ */
 export async function* readDataStream(
   reader: ReadableStreamDefaultReader<Uint8Array>,
   {
diff --git a/packages/core/shared/stream-parts.ts b/packages/core/shared/stream-parts.ts
@@ -328,11 +328,11 @@ export const StreamStringPrefixes = {
 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.
+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 firstSeparatorIndex = line.indexOf(':');
@@ -356,10 +356,10 @@ export const parseStreamPart = (line: string): StreamPartType => {
 };
 
 /**
- * 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.
+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,
diff --git a/packages/core/shared/utils.ts b/packages/core/shared/utils.ts
@@ -9,6 +9,12 @@ export * from './generate-id';
 // TODO remove (breaking change)
 export { generateId as nanoid } from './generate-id';
 
+// Export stream data utilities for custom stream implementations,
+// both on the client and server side.
+export type { StreamPart } from './stream-parts';
+export { formatStreamPart, parseStreamPart } from './stream-parts';
+export { readDataStream } from './read-data-stream';
+
 // simple decoder signatures:
 function createChunkDecoder(): (chunk: Uint8Array | undefined) => string;
 function createChunkDecoder(

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +'ai': patch
 +---
++
 +Expose formatStreamPart, parseStreamPart, and readDataStream helpers.