ai-stream-utils 0.0.3 → 1.4.2

package/README.md

@@ -30,69 +30,32 @@ npm install ai-stream-utils
 
 ## Overview
 
-| Function | Object | Use Case |
-|----------|-------------|----------|
-| [`mapUIMessageStream`](#mapuimessagestream) | [UIMessageChunk](https://github.com/vercel/ai/blob/main/packages/ai/src/ui-message-stream/ui-message-chunks.ts) | Transform chunks while streaming to the client |
-| [`flatMapUIMessageStream`](#flatmapuimessagestream) | [UIMessagePart](https://ai-sdk.dev/docs/reference/ai-sdk-core/ui-message#uimessagepart-types) | Buffer chunks until a part is complete, then transform the part |
-| [`filterUIMessageStream`](#filteruimessagestream) | [UIMessageChunk](https://github.com/vercel/ai/blob/main/packages/ai/src/ui-message-stream/ui-message-chunks.ts) | Filter chunks while streaming to the client |
+| Function | Input | Returns | Use Case |
+|----------|-------------|---------|----------|
+| [`mapUIMessageStream`](#mapuimessagestream) | [UIMessageChunk](https://github.com/vercel/ai/blob/main/packages/ai/src/ui-message-stream/ui-message-chunks.ts) | `chunk \| chunk[] \| null` | Transform or filter chunks in real-time (e.g., smooth streaming) |
+| [`flatMapUIMessageStream`](#flatmapuimessagestream) | [UIMessagePart](https://ai-sdk.dev/docs/reference/ai-sdk-core/ui-message#uimessagepart-types) | `part \| part[] \| null` | Buffer until complete, then transform (e.g., redact tool output) |
+| [`filterUIMessageStream`](#filteruimessagestream) | [UIMessageChunk](https://github.com/vercel/ai/blob/main/packages/ai/src/ui-message-stream/ui-message-chunks.ts) | `boolean` | Include/exclude parts by type (e.g., hide reasoning) |
 
 ## Usage
 
 ### `mapUIMessageStream`
 
-Transform or filter individual chunks as they stream through. The map function receives the chunk and a partial representation of the part it belongs to.
+The `mapUIMessageStream` function operates on chunks and can be used to transform or filter individual chunks as they stream through. It receives the current chunk and the partial part representing all already processed chunks.
 
 ```typescript
 import { mapUIMessageStream } from 'ai-stream-utils';
-import { streamText } from 'ai';
 
-const tools = {
-  weather: tool({
-    description: 'Get the weather in a location',
-    inputSchema: z.object({
-      location: z.string().describe('The location to get the weather for'),
-    }),
-    execute: ({ location }) => ({
-      location,
-      temperature: 72 + Math.floor(Math.random() * 21) - 10,
-      unit: "C",
-    }),
-  }),
-};
-
-const result = streamText({
-  model,
-  prompt: 'What is the weather in Tokyo?',
-  tools,
-});
-
-// Filter out tool-call chunks by part type
-const stream = mapUIMessageStream(
-  result.toUIMessageStream<MyUIMessage>(),
-  ({ chunk, part }) => {
-    if (part.type === "tool-weather") {
-      return null;
-    }
-    return chunk;
-  }
-);
-
-// Transform text chunks to uppercase
 const stream = mapUIMessageStream(
   result.toUIMessageStream<MyUIMessage>(),
   ({ chunk, part }) => {
+    // Transform: modify the chunk
     if (chunk.type === 'text-delta') {
       return { ...chunk, delta: chunk.delta.toUpperCase() };
     }
-    return chunk;
-  }
-);
-
-// Access chunk history and index
-const stream = mapUIMessageStream(
-  result.toUIMessageStream<MyUIMessage>(),
-  ({ chunk }, { index, chunks }) => {
-    console.log(`Chunk ${index}, total seen: ${chunks.length}`);
+    // Filter: return null to exclude chunks
+    if (part.type === 'tool-weather') {
+      return null;
+    }
     return chunk;
   }
 );
@@ -100,49 +63,21 @@ const stream = mapUIMessageStream(
 
 ### `flatMapUIMessageStream`
 
-Buffer all chunks for a part until it's complete, then transform the complete part. This is useful when you need access to the full part content before deciding how to transform it.
-
-When a predicate is provided (e.g., `partTypeIs('text')`), only matching parts are buffered for transformation. Non-matching parts stream through immediately without buffering, preserving real-time streaming behavior.
+The `flatMapUIMessageStream` function operates on parts. It buffers all chunks of a particular type (e.g. text parts) until the part is complete and then transforms or filters the complete part. The optional predicate `partTypeIs()` can be used to selectively buffer only specific parts while streaming others through immediately.
 
 ```typescript
 import { flatMapUIMessageStream, partTypeIs } from 'ai-stream-utils';
 
-// Filter out reasoning parts
-const stream = flatMapUIMessageStream(
-  result.toUIMessageStream<MyUIMessage>(),
-  ({ part }) => part.type === 'reasoning' ? null : part
-);
-
-// Transform text content
 const stream = flatMapUIMessageStream(
   result.toUIMessageStream<MyUIMessage>(),
+  // Predicate to only buffer tool-weather parts and pass through other parts
+  partTypeIs('tool-weather'),
   ({ part }) => {
-    if (part.type === 'text') {
-      return { ...part, text: part.text.toUpperCase() };
+    // Transform: modify the complete part
+    if (part.state === 'output-available') {
+      return { ...part, output: { ...part.output, temperature: toFahrenheit(part.output.temperature) } };
     }
-    return part;
-  }
-);
-
-// Buffer only specific parts, pass through others immediately
-const stream = flatMapUIMessageStream(
-  result.toUIMessageStream<MyUIMessage>(),
-  partTypeIs('text'),
-  ({ part }) => ({ ...part, text: part.text.toUpperCase() })
-);
-
-// Buffer multiple part types
-const stream = flatMapUIMessageStream(
-  result.toUIMessageStream<MyUIMessage>(),
-  partTypeIs(['text', 'reasoning']),
-  ({ part }) => part // part is typed as TextUIPart | ReasoningUIPart
-);
-
-// Access part history
-const stream = flatMapUIMessageStream(
-  result.toUIMessageStream<MyUIMessage>(),
-  ({ part }, { index, parts }) => {
-    console.log(`Part ${index}, previous parts:`, parts.slice(0, -1));
+    // Filter: return null to exclude parts
     return part;
   }
 );
@@ -150,7 +85,7 @@ const stream = flatMapUIMessageStream(
 
 ### `filterUIMessageStream`
 
-Filter individual chunks as they stream through. This is a convenience wrapper around `mapUIMessageStream` that provides a simpler API for filtering chunks by part type. Use the `includeParts()` and `excludeParts()` helper functions for common patterns, or provide a custom filter function.
+The `filterUIMessageStream` function is a convenience function around `mapUIMessageStream` with a simpler API to filter chunks by part type. It provides the  `includeParts()` and `excludeParts()` predicates for common patterns.
 
 ```typescript
 import { filterUIMessageStream, includeParts, excludeParts } from 'ai-stream-utils';
@@ -170,16 +105,280 @@ const stream = filterUIMessageStream(
 // Custom filter function
 const stream = filterUIMessageStream(
   result.toUIMessageStream<MyUIMessage>(),
-  ({ part }, { index }) => {
-    // Include text parts
+  ({ part, chunk }) => {
     if (part.type === 'text') return true;
-    // Include only first 5 parts
-    if (index < 5) return true;
+    if (chunk.type === 'tool-input-available') return true;
     return false;
   }
 );
 ```
 
+## Examples
+
+### Smooth Streaming
+
+Buffers multiple text chunks into a string, splits at word boundaries and re-emits each word as a separate chunk for smoother UI rendering. See [examples/smooth-streaming.ts](./examples/smooth-streaming.ts) for the full implementation.
+
+```typescript
+import { mapUIMessageStream } from 'ai-stream-utils';
+
+const WORD_REGEX = /\S+\s+/m;
+let buffer = '';
+
+const smoothedStream = mapUIMessageStream(
+  result.toUIMessageStream(),
+  ({ chunk }) => {
+    if (chunk.type !== 'text-delta') {
+      // Flush buffer on non-text chunks
+      if (buffer.length > 0) {
+        const flushed = { type: 'text-delta' as const, id: chunk.id, delta: buffer };
+        buffer = '';
+        return [flushed, chunk];
+      }
+      return chunk;
+    }
+
+    // Append the text delta to the buffer
+    buffer += chunk.delta;
+    const chunks = [];
+    
+    let match;
+    while ((match = WORD_REGEX.exec(buffer)) !== null) {
+      chunks.push({ type: 'text-delta', id: chunk.id, delta: buffer.slice(0, match.index + match[0].length) });
+      buffer = buffer.slice(match.index + match[0].length);
+    }
+    // Emit the word-by-word chunks
+    return chunks;
+  }
+);
+
+// Output: word-by-word streaming
+// { type: 'text-delta', delta: 'Why ' }
+// { type: 'text-delta', delta: "don't " }
+// { type: 'text-delta', delta: 'scientists ' }
+```
+
+### Redacting Sensitive Data
+
+Buffer tool calls until complete, then redact sensitive fields before streaming to the client. See [examples/order-lookup.ts](./examples/order-lookup.ts) for the full example.
+
+```typescript
+import { flatMapUIMessageStream, partTypeIs } from 'ai-stream-utils';
+
+const tools = {
+  lookupOrder: tool({
+    description: 'Look up order details by order ID',
+    inputSchema: z.object({
+      orderId: z.string().describe('The order ID to look up'),
+    }),
+    execute: ({ orderId }) => ({
+      orderId,
+      status: 'shipped',
+      items: ['iPhone 15'],
+      total: 1299.99,
+      email: 'customer@example.com',        // Sensitive
+      address: '123 Main St, SF, CA 94102', // Sensitive
+    }),
+  }),
+};
+
+const result = streamText({
+  model: openai('gpt-4o'),
+  prompt: 'Where is my order #12345?',
+  tools,
+});
+
+// Buffer tool-lookupOrder parts, stream text parts immediately
+const redactedStream = flatMapUIMessageStream(
+  result.toUIMessageStream<MyUIMessage>(),
+  partTypeIs('tool-lookupOrder'),
+  ({ part }) => {
+    if (part.state === 'output-available') {
+      return {
+        ...part,
+        output: {
+          ...part.output,
+          email: '[REDACTED]',
+          address: '[REDACTED]',
+        },
+      };
+    }
+    return part;
+  },
+);
+
+// Text streams immediately, tool output is redacted:
+// { type: 'text-delta', delta: 'Let me look that up...' }
+// { type: 'tool-output-available', output: { orderId: '12345', email: '[REDACTED]', address: '[REDACTED]' } }
+```
+
+### Conditional Part Injection
+
+Inspect previously streamed parts to conditionally inject new parts. This example creates a text part from a tool call message if the model didn't generate one. See [examples/ask-permission.ts](./examples/ask-permission.ts) for the full example.
+
+```typescript
+import { flatMapUIMessageStream, partTypeIs } from 'ai-stream-utils';
+
+const tools = {
+  askForPermission: tool({
+    description: 'Ask for permission to access current location',
+    inputSchema: z.object({
+      message: z.string().describe('The message to ask for permission'),
+    }),
+  }),
+};
+
+const result = streamText({
+  model: openai('gpt-4o'),
+  prompt: 'Is it sunny today?',
+  tools,
+});
+
+// Buffer askForPermission tool calls, check if text was already generated
+const stream = flatMapUIMessageStream(
+  result.toUIMessageStream<MyUIMessage>(),
+  partTypeIs('tool-askForPermission'),
+  (current, context) => {
+    if (current.part.state === 'input-available') {
+      // Check if a text part was already streamed
+      const hasTextPart = context.parts.some((p) => p.type === 'text');
+      
+      if (!hasTextPart) {
+        // Inject a text part from the tool call message
+        return [
+          { type: 'text', text: current.part.input.message },
+          current.part,
+        ];
+      }
+    }
+    return current.part;
+  },
+);
+
+// If model only generated tool call, we inject the text:
+// { type: 'text', text: 'May I access your location?' }
+// { type: 'tool-askForPermission', input: { message: 'May I access your location?' } }
+```
+
+### Transform Tool Output
+
+Transform tool outputs on-the-fly, such as converting temperature units. See [examples/weather.ts](./examples/weather.ts) for the full example.
+
+```typescript
+import { flatMapUIMessageStream, partTypeIs } from 'ai-stream-utils';
+
+const toFahrenheit = (celsius: number) => (celsius * 9) / 5 + 32;
+
+const tools = {
+  weather: tool({
+    description: 'Get the weather in a location',
+    inputSchema: z.object({ location: z.string() }),
+    execute: ({ location }) => ({
+      location,
+      temperature: 22, // Celsius from API
+      unit: 'C',
+    }),
+  }),
+};
+
+const result = streamText({
+  model: openai('gpt-4o'),
+  prompt: 'What is the weather in Tokyo?',
+  tools,
+});
+
+// Convert Celsius to Fahrenheit before streaming to client
+const stream = flatMapUIMessageStream(
+  result.toUIMessageStream<MyUIMessage>(),
+  partTypeIs('tool-weather'),
+  ({ part }) => {
+    if (part.state === 'output-available') {
+      return {
+        ...part,
+        output: {
+          ...part.output,
+          temperature: toFahrenheit(part.output.temperature),
+          unit: 'F',
+        },
+      };
+    }
+    return part;
+  },
+);
+
+// Output is converted:
+// { type: 'tool-output-available', output: { location: 'Tokyo', temperature: 71.6, unit: 'F' } }
+```
+
+## Stream Utilities
+
+Helper functions for converting between streams, arrays, and async iterables.
+
+| Function | Converts | To |
+|----------|----------|-----|
+| `createAsyncIterableStream` | `ReadableStream<T>` | `AsyncIterableStream<T>` |
+| `convertArrayToStream` | `Array<T>` | `ReadableStream<T>` |
+| `convertAsyncIterableToStream` | `AsyncIterable<T>` | `ReadableStream<T>` |
+| `convertAsyncIterableToArray` | `AsyncIterable<T>` | `Promise<Array<T>>` |
+| `convertStreamToArray` | `ReadableStream<T>` | `Promise<Array<T>>` |
+
+### `createAsyncIterableStream`
+
+Adds async iterator protocol to a `ReadableStream`, enabling `for await...of` loops.
+
+```typescript
+import { createAsyncIterableStream } from 'ai-stream-utils/utils';
+
+const asyncStream = createAsyncIterableStream(readableStream);
+for await (const chunk of asyncStream) {
+  console.log(chunk);
+}
+```
+
+### `convertArrayToStream`
+
+Converts an array to a `ReadableStream` that emits each element.
+
+```typescript
+import { convertArrayToStream } from 'ai-stream-utils/utils';
+
+const stream = convertArrayToStream([1, 2, 3]);
+```
+
+### `convertAsyncIterableToStream`
+
+Converts an async iterable (e.g., async generator) to a `ReadableStream`.
+
+```typescript
+import { convertAsyncIterableToStream } from 'ai-stream-utils/utils';
+
+async function* generator() {
+  yield 1;
+  yield 2;
+}
+const stream = convertAsyncIterableToStream(generator());
+```
+
+### `convertAsyncIterableToArray`
+
+Collects all values from an async iterable into an array.
+
+```typescript
+import { convertAsyncIterableToArray } from 'ai-stream-utils/utils';
+
+const array = await convertAsyncIterableToArray(asyncIterable);
+```
+
+### `convertStreamToArray`
+
+Consumes a `ReadableStream` and collects all chunks into an array.
+
+```typescript
+import { convertStreamToArray } from 'ai-stream-utils/utils';
+
+const array = await convertStreamToArray(readableStream);
+```
+
 ## Type Safety
 
 The [`toUIMessageStream()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/stream-text#to-ui-message-stream) from [`streamText()`](https://ai-sdk.dev/docs/reference/ai-sdk-core/stream-text) returns a generic `ReadableStream<UIMessageChunk>`, which means the part types cannot be inferred automatically.
@@ -272,17 +471,11 @@ function mapUIMessageStream<UI_MESSAGE extends UIMessage>(
 
 type MapUIMessageStreamFn<UI_MESSAGE extends UIMessage> = (
   input: MapInput<UI_MESSAGE>,
-  context: MapContext<UI_MESSAGE>,
-) => InferUIMessageChunk<UI_MESSAGE> | null;
+) => InferUIMessageChunk<UI_MESSAGE> | InferUIMessageChunk<UI_MESSAGE>[] | null;
 
 type MapInput<UI_MESSAGE extends UIMessage> = {
   chunk: InferUIMessageChunk<UI_MESSAGE>;
-  part: InferPartialUIMessagePart<UI_MESSAGE>;
-};
-
-type MapContext<UI_MESSAGE extends UIMessage> = {
-  index: number;
-  chunks: InferUIMessageChunk<UI_MESSAGE>[];
+  part: InferUIMessagePart<UI_MESSAGE>;
 };
 ```
 
@@ -305,7 +498,7 @@ function flatMapUIMessageStream<UI_MESSAGE extends UIMessage, PART extends Infer
 type FlatMapUIMessageStreamFn<UI_MESSAGE extends UIMessage, PART = InferUIMessagePart<UI_MESSAGE>> = (
   input: FlatMapInput<UI_MESSAGE, PART>,
   context: FlatMapContext<UI_MESSAGE>,
-) => PART | null;
+) => InferUIMessagePart<UI_MESSAGE> | InferUIMessagePart<UI_MESSAGE>[] | null;
 
 type FlatMapInput<UI_MESSAGE extends UIMessage, PART = InferUIMessagePart<UI_MESSAGE>> = {
   part: PART;
@@ -325,7 +518,7 @@ function partTypeIs<UI_MESSAGE extends UIMessage, T extends InferUIMessagePartTy
 ): FlatMapUIMessageStreamPredicate<UI_MESSAGE, Extract<InferUIMessagePart<UI_MESSAGE>, { type: T }>>
 
 type FlatMapUIMessageStreamPredicate<UI_MESSAGE extends UIMessage, PART extends InferUIMessagePart<UI_MESSAGE>> = 
-  (part: InferPartialUIMessagePart<UI_MESSAGE>) => boolean;
+  (part: InferUIMessagePart<UI_MESSAGE>) => boolean;
 ```
 
 ### `filterUIMessageStream`

package/dist/create-async-iterable-stream-x_DKVIDi.mjs

@@ -0,0 +1,59 @@
+//#region src/utils/create-async-iterable-stream.ts
+/**
+* Converts a ReadableStream to an AsyncIterableStream.
+* Copied from https://github.com/vercel/ai/blob/main/packages/ai/src/util/async-iterable-stream.ts
+*/
+function createAsyncIterableStream(source) {
+	/** Pipe through a TransformStream to ensure a fresh, unlocked stream. */
+	const stream = source.pipeThrough(new TransformStream());
+	/** Implements the async iterator protocol for the stream. */
+	return Object.assign(stream, { [Symbol.asyncIterator]() {
+		const reader = stream.getReader();
+		let finished = false;
+		/** Cleans up the reader by cancelling and releasing the lock. */
+		async function cleanup(cancelStream) {
+			finished = true;
+			try {
+				if (cancelStream) await reader.cancel?.();
+			} finally {
+				try {
+					reader.releaseLock();
+				} catch {}
+			}
+		}
+		return {
+			async next() {
+				if (finished) return {
+					done: true,
+					value: void 0
+				};
+				const { done, value } = await reader.read();
+				if (done) {
+					await cleanup(true);
+					return {
+						done: true,
+						value: void 0
+					};
+				}
+				return {
+					done: false,
+					value
+				};
+			},
+			async return() {
+				await cleanup(true);
+				return {
+					done: true,
+					value: void 0
+				};
+			},
+			async throw(err) {
+				await cleanup(true);
+				throw err;
+			}
+		};
+	} });
+}
+
+//#endregion
+export { createAsyncIterableStream as t };