feat (ai/core): support multiple stream text transforms (#4387)

vercel · Jan 14, 2025 · 3491f78 · 3491f78
1 parent cd06d5f
commit 3491f78
Show file tree

Hide file tree

Showing 7 changed files with 483 additions and 381 deletions.
diff --git a/.changeset/cool-rings-tell.md b/.changeset/cool-rings-tell.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/core): support multiple stream text transforms
diff --git a/content/docs/03-ai-sdk-core/05-generating-text.mdx b/content/docs/03-ai-sdk-core/05-generating-text.mdx
@@ -319,6 +319,18 @@ const stopWordTransform =
     });
 ```
 
+#### Multiple transformations
+
+You can also provide multiple transformations. They are applied in the order they are provided.
+
+```tsx highlight="4"
+const result = streamText({
+  model,
+  prompt,
+  experimental_transform: [firstTransform, secondTransform],
+});
+```
+
 ## Generating Long Text
 
 Most language models have an output limit that is much shorter than their context window.

diff --git a/content/docs/07-reference/01-ai-sdk-core/02-stream-text.mdx b/content/docs/07-reference/01-ai-sdk-core/02-stream-text.mdx
@@ -471,23 +471,35 @@ To see `streamText` in action, check out [these examples](#examples).
     },
     {
       name: 'experimental_transform',
-      type: '(options: TransformOptions) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>',
+      type: 'StreamTextTransform | Array<StreamTextTransform>',
       isOptional: true,
       description:
-        'Optional transformation that is applied to the stream. Tools are passed in for type inference.',
+        'Optional stream transformations. They are applied in the order they are provided. The stream transformations must maintain the stream structure for streamText to work correctly.',
       properties: [
         {
-          type: 'TransformOptions',
+          type: 'StreamTextTransform',
           parameters: [
             {
-              name: 'stopStream',
-              type: '() => void',
-              description: 'A function that stops the stream.',
-            },
-            {
-              name: 'tools',
-              type: 'TOOLS',
-              description: 'The tools that are available.',
+              name: 'transform',
+              type: '(options: TransformOptions) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>',
+              description: 'A transformation that is applied to the stream.',
+              properties: [
+                {
+                  type: 'TransformOptions',
+                  parameters: [
+                    {
+                      name: 'stopStream',
+                      type: '() => void',
+                      description: 'A function that stops the stream.',
+                    },
+                    {
+                      name: 'tools',
+                      type: 'TOOLS',
+                      description: 'The tools that are available.',
+                    },
+                  ],
+                },
+              ],
             },
           ],
         },

diff --git a/packages/ai/core/generate-text/__snapshots__/stream-text.test.ts.snap b/packages/ai/core/generate-text/__snapshots__/stream-text.test.ts.snap
@@ -2181,6 +2181,42 @@ exports[`streamText > options.transform > with base transformation > telemetry s
 ]
 `;
 
+exports[`streamText > options.transform > with transformation that aborts stream > options.onStepFinish should be called 1`] = `
+{
+  "experimental_providerMetadata": undefined,
+  "finishReason": "stop",
+  "isContinued": false,
+  "logprobs": undefined,
+  "request": {},
+  "response": {
+    "id": "response-id",
+    "messages": [
+      {
+        "content": [
+          {
+            "text": "Hello, ",
+            "type": "text",
+          },
+        ],
+        "role": "assistant",
+      },
+    ],
+    "modelId": "mock-model-id",
+    "timestamp": 1970-01-01T00:00:00.000Z,
+  },
+  "stepType": "initial",
+  "text": "Hello, ",
+  "toolCalls": [],
+  "toolResults": [],
+  "usage": {
+    "completionTokens": NaN,
+    "promptTokens": NaN,
+    "totalTokens": NaN,
+  },
+  "warnings": [],
+}
+`;
+
 exports[`streamText > result.fullStream > should filter out empty text deltas 1`] = `
 [
   {
@@ -3315,39 +3351,3 @@ exports[`streamText > tools with custom schema > should send tool calls 1`] = `
   },
 ]
 `;
-
-exports[`streamText > with transformation that aborts stream > options.onStepFinish should be called 1`] = `
-{
-  "experimental_providerMetadata": undefined,
-  "finishReason": "stop",
-  "isContinued": false,
-  "logprobs": undefined,
-  "request": {},
-  "response": {
-    "id": "response-id",
-    "messages": [
-      {
-        "content": [
-          {
-            "text": "Hello, ",
-            "type": "text",
-          },
-        ],
-        "role": "assistant",
-      },
-    ],
-    "modelId": "mock-model-id",
-    "timestamp": 1970-01-01T00:00:00.000Z,
-  },
-  "stepType": "initial",
-  "text": "Hello, ",
-  "toolCalls": [],
-  "toolResults": [],
-  "usage": {
-    "completionTokens": NaN,
-    "promptTokens": NaN,
-    "totalTokens": NaN,
-  },
-  "warnings": [],
-}
-`;
diff --git a/packages/ai/core/generate-text/index.ts b/packages/ai/core/generate-text/index.ts
@@ -1,11 +1,12 @@
 export { generateText } from './generate-text';
 export type { GenerateTextResult } from './generate-text-result';
 export * as Output from './output';
+export { smoothStream } from './smooth-stream';
 export type { StepResult } from './step-result';
 export { streamText } from './stream-text';
+export type { StreamTextTransform } from './stream-text';
 export type { StreamTextResult, TextStreamPart } from './stream-text-result';
 export type { ToolCallRepairFunction } from './tool-call-repair';
-export { smoothStream } from './smooth-stream';
 
 // TODO 4.1: rename to ToolCall and ToolResult, deprecate old names
 export type {