fix (ai/core): change streamText.experimental_transform signature to …

…support tool type inference (#4090)
vercel · Dec 13, 2024 · 54bbf21 · 54bbf21
1 parent 5b59fce
commit 54bbf21
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 44 deletions.
diff --git a/.changeset/two-mirrors-fry.md b/.changeset/two-mirrors-fry.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ai/core): change streamText.experimental_transform signature to support tool type inference
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,9 +471,10 @@ To see `streamText` in action, check out [these examples](#examples).
     },
     {
       name: 'experimental_transform',
-      type: 'TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>',
+      type: '({ tools: TOOLS }) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>',
       isOptional: true,
-      description: 'Optional transformation that is applied to the stream.',
+      description:
+        'Optional transformation that is applied to the stream. Tools are passed in for type inference.',
     },
     {
       name: 'experimental_providerMetadata',

diff --git a/packages/ai/core/generate-text/smooth-stream.test.ts b/packages/ai/core/generate-text/smooth-stream.test.ts
@@ -21,7 +21,7 @@ describe('smoothStream', () => {
             return Promise.resolve();
           },
         },
-      }),
+      })({ tools: {} }),
     );
 
     // Get a reader and read chunks
@@ -70,7 +70,7 @@ describe('smoothStream', () => {
             return Promise.resolve();
           },
         },
-      }),
+      })({ tools: {} }),
     );
 
     // Get a reader and read chunks

diff --git a/packages/ai/core/generate-text/smooth-stream.ts b/packages/ai/core/generate-text/smooth-stream.ts
@@ -20,39 +20,42 @@ export function smoothStream<TOOLS extends Record<string, CoreTool>>({
   _internal?: {
     delay?: (delayInMs: number) => Promise<void>;
   };
-} = {}): TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> {
+} = {}): (options: {
+  tools: TOOLS;
+}) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> {
   let buffer = '';
 
-  return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
-    async transform(chunk, controller) {
-      if (chunk.type === 'step-finish') {
-        if (buffer.length > 0) {
-          controller.enqueue({ type: 'text-delta', textDelta: buffer });
-          buffer = '';
+  return () =>
+    new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
+      async transform(chunk, controller) {
+        if (chunk.type === 'step-finish') {
+          if (buffer.length > 0) {
+            controller.enqueue({ type: 'text-delta', textDelta: buffer });
+            buffer = '';
+          }
+
+          controller.enqueue(chunk);
+          return;
         }
 
-        controller.enqueue(chunk);
-        return;
-      }
-
-      if (chunk.type !== 'text-delta') {
-        controller.enqueue(chunk);
-        return;
-      }
+        if (chunk.type !== 'text-delta') {
+          controller.enqueue(chunk);
+          return;
+        }
 
-      buffer += chunk.textDelta;
+        buffer += chunk.textDelta;
 
-      // Stream out complete words when whitespace is found
-      while (buffer.match(/\s/)) {
-        const whitespaceIndex = buffer.search(/\s/);
-        const word = buffer.slice(0, whitespaceIndex + 1);
-        controller.enqueue({ type: 'text-delta', textDelta: word });
-        buffer = buffer.slice(whitespaceIndex + 1);
+        // Stream out complete words when whitespace is found
+        while (buffer.match(/\s/)) {
+          const whitespaceIndex = buffer.search(/\s/);
+          const word = buffer.slice(0, whitespaceIndex + 1);
+          controller.enqueue({ type: 'text-delta', textDelta: word });
+          buffer = buffer.slice(whitespaceIndex + 1);
 
-        if (delayInMs > 0) {
-          await delay(delayInMs);
+          if (delayInMs > 0) {
+            await delay(delayInMs);
+          }
         }
-      }
-    },
-  });
+      },
+    });
 }
diff --git a/packages/ai/core/generate-text/stream-text.test.ts b/packages/ai/core/generate-text/stream-text.test.ts
@@ -3277,14 +3277,15 @@ describe('streamText', () => {
             rawCall: { rawPrompt: 'prompt', rawSettings: {} },
           }),
         }),
-        experimental_transform: new TransformStream({
-          transform(chunk, controller) {
-            if (chunk.type === 'text-delta') {
-              chunk.textDelta = chunk.textDelta.toUpperCase();
-            }
-            controller.enqueue(chunk);
-          },
-        }),
+        experimental_transform: () =>
+          new TransformStream({
+            transform(chunk, controller) {
+              if (chunk.type === 'text-delta') {
+                chunk.textDelta = chunk.textDelta.toUpperCase();
+              }
+              controller.enqueue(chunk);
+            },
+          }),
         prompt: 'test-input',
       });
 

diff --git a/packages/ai/core/generate-text/stream-text.ts b/packages/ai/core/generate-text/stream-text.ts
@@ -190,10 +190,9 @@ Enable streaming of tool call deltas as they are generated. Disabled by default.
     /**
 Optional transformation that is applied to the stream.
      */
-    experimental_transform?: TransformStream<
-      TextStreamPart<TOOLS>,
-      TextStreamPart<TOOLS>
-    >;
+    experimental_transform?: (options: {
+      tools: TOOLS; // for type inference
+    }) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>;
 
     /**
 Callback that is called for each chunk of the stream. The stream processing will pause until the callback promise is resolved.
@@ -254,7 +253,7 @@ Details for all steps.
     tools,
     toolChoice,
     toolCallStreaming,
-    transform,
+    transform: transform?.({ tools: tools as TOOLS }),
     activeTools,
     repairToolCall,
     maxSteps,