feat (ai/streams): add StreamData support to streamToResponse (#1615)

Co-authored-by: Valtteri Karesto <valtteri.karesto@gmail.com>

Author: lgrammel

.changeset/unlucky-drinks-fry.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/streams): add StreamData support to streamToResponse

content/docs/07-reference/stream-helpers/05-stream-to-response.mdx

@@ -5,12 +5,64 @@ description: Learn to use streamToResponse helper function in your application.
 
 # `streamToResponse`
 
-This method will pipe a ReadableStream to a Node.js ServerResponse object. It can be helpful to combine this with other AI stream utilities, such as OpenAIStream, in Node.js environments. Similar to StreamingTextResponse, it automatically sets the status code to 200 and the Content-Type header to 'text/plain; charset=utf-8'
+`streamToResponse` pipes an AI stream to a Node.js `ServerResponse` object and sets the status code and headers.
+
+This is useful to create AI stream responses in environments that use `ServerResponse` objects, such as Node.js HTTP servers.
+
+The status code and headers can be configured using the `options` parameter.
+By default, the status code is set to 200 and the Content-Type header is set to `text/plain; charset=utf-8`.
 
 ## Import
 
 ### React
 
 <Snippet text={`import { streamToResponse } from "ai"`} prompt={false} />
 
-<ReferenceTable packageName="streams" functionName="streamToResponse" />
+<ObjectTableList
+  content={{
+    Parameters: [
+      {
+        name: 'stream',
+        type: 'ReadableStream',
+        description:
+          'The Web Stream to pipe to the response. It can be the return value of OpenAIStream, HuggingFaceStream, AnthropicStream, or an AIStream instance.',
+      },
+      {
+        name: 'response',
+        type: 'ServerResponse',
+        description:
+          'The Node.js ServerResponse object to pipe the stream to. This is usually the second argument of a Node.js HTTP request handler.',
+      },
+      {
+        name: 'options',
+        type: 'Options',
+        description: 'Configure the response',
+        properties: [
+          {
+            type: 'Options',
+            parameters: [
+              {
+                name: 'status',
+                type: 'number',
+                description:
+                  'The status code to set on the response. Defaults to `200`.',
+              },
+              {
+                name: 'headers',
+                type: 'Record<string, string>',
+                description:
+                  "Additional headers to set on the response. Defaults to `{ 'Content-Type': 'text/plain; charset=utf-8' }`.",
+              },
+            ],
+          },
+        ],
+      },
+      {
+        name: 'data',
+        type: 'StreamData',
+        description:
+          'StreamData object for forwarding additional data to the client.',
+      },
+    ],
+  }}
+/>

packages/core/streams/index.ts

@@ -15,5 +15,6 @@ export * from './mistral-stream';
 export * from './openai-stream';
 export * from './replicate-stream';
 export * from './stream-data';
+export * from './stream-to-response';
 export * from './streaming-react-response';
 export * from './streaming-text-response';

packages/core/streams/stream-to-response.ts

@@ -0,0 +1,37 @@
+import type { ServerResponse } from 'node:http';
+import { StreamData } from './stream-data';
+import { mergeStreams } from '../core/util/merge-streams';
+
+/**
+ * A utility function to stream a ReadableStream to a Node.js response-like object.
+ */
+export function streamToResponse(
+  res: ReadableStream,
+  response: ServerResponse,
+  init?: { headers?: Record<string, string>; status?: number },
+  data?: StreamData,
+) {
+  response.writeHead(init?.status ?? 200, {
+    'Content-Type': 'text/plain; charset=utf-8',
+    ...init?.headers,
+  });
+
+  let processedStream = res;
+
+  if (data) {
+    processedStream = mergeStreams(data.stream, res);
+  }
+
+  const reader = processedStream.getReader();
+  function read() {
+    reader.read().then(({ done, value }: { done: boolean; value?: any }) => {
+      if (done) {
+        response.end();
+        return;
+      }
+      response.write(value);
+      read();
+    });
+  }
+  read();
+}

packages/core/streams/streaming-text-response.ts

@@ -1,6 +1,5 @@
-import type { ServerResponse } from 'node:http';
-import { StreamData } from './stream-data';
 import { mergeStreams } from '../core/util/merge-streams';
+import { StreamData } from './stream-data';
 
 /**
  * A utility class for streaming text responses.
@@ -23,30 +22,3 @@ export class StreamingTextResponse extends Response {
     });
   }
 }
-
-/**
- * A utility function to stream a ReadableStream to a Node.js response-like object.
- */
-export function streamToResponse(
-  res: ReadableStream,
-  response: ServerResponse,
-  init?: { headers?: Record<string, string>; status?: number },
-) {
-  response.writeHead(init?.status || 200, {
-    'Content-Type': 'text/plain; charset=utf-8',
-    ...init?.headers,
-  });
-
-  const reader = res.getReader();
-  function read() {
-    reader.read().then(({ done, value }: { done: boolean; value?: any }) => {
-      if (done) {
-        response.end();
-        return;
-      }
-      response.write(value);
-      read();
-    });
-  }
-  read();
-}