feat (ai/core): consolidate whitespace in smooth stream (#4244)

vercel · Jan 2, 2025 · bc4cd19 · bc4cd19
1 parent 647f39d
commit bc4cd19
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 6 deletions.
diff --git a/.changeset/many-colts-clean.md b/.changeset/many-colts-clean.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/core): consolidate whitespace in smooth stream
diff --git a/packages/ai/core/generate-text/smooth-stream.test.ts b/packages/ai/core/generate-text/smooth-stream.test.ts
@@ -124,4 +124,70 @@ describe('smoothStream', () => {
       },
     ]);
   });
+
+  it('should keep longer whitespace sequences together', async () => {
+    const events: any[] = [];
+
+    const stream = convertArrayToReadableStream([
+      { textDelta: 'First line', type: 'text-delta' },
+      { textDelta: ' \n\n', type: 'text-delta' },
+      { textDelta: '  ', type: 'text-delta' },
+      { textDelta: '  Multiple spaces', type: 'text-delta' },
+      { textDelta: '\n    Indented', type: 'text-delta' },
+      { type: 'step-finish' },
+      { type: 'finish' },
+    ]).pipeThrough(
+      smoothStream({
+        delayInMs: 10,
+        _internal: {
+          delay: () => {
+            events.push('delay');
+            return Promise.resolve();
+          },
+        },
+      })({ tools: {} }),
+    );
+
+    const reader = stream.getReader();
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      events.push(value);
+    }
+
+    expect(events).toEqual([
+      'delay',
+      {
+        textDelta: 'First ',
+        type: 'text-delta',
+      },
+      'delay',
+      {
+        textDelta: 'line \n\n',
+        type: 'text-delta',
+      },
+      'delay',
+      {
+        // note: leading whitespace is included here
+        // because it is part of the new chunk:
+        textDelta: '    Multiple ',
+        type: 'text-delta',
+      },
+      'delay',
+      {
+        textDelta: 'spaces\n    ',
+        type: 'text-delta',
+      },
+      {
+        textDelta: 'Indented',
+        type: 'text-delta',
+      },
+      {
+        type: 'step-finish',
+      },
+      {
+        type: 'finish',
+      },
+    ]);
+  });
 });
diff --git a/packages/ai/core/generate-text/smooth-stream.ts b/packages/ai/core/generate-text/smooth-stream.ts
@@ -45,12 +45,13 @@ export function smoothStream<TOOLS extends Record<string, CoreTool>>({
 
         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 including their optional leading
+        // and required trailing whitespace sequences
+        const regexp = /\s*\S+\s+/m;
+        while (regexp.test(buffer)) {
+          const chunk = buffer.match(regexp)![0];
+          controller.enqueue({ type: 'text-delta', textDelta: chunk });
+          buffer = buffer.slice(chunk.length);
 
           if (delayInMs > 0) {
             await delay(delayInMs);