ai-filter-stream

This library allows you filter UI message chunks returned from streamText() by their corresponding UI message part type.

Why?

The AI SDK UI message stream created by toUIMessageStream() streams all parts (text, tools, data, etc.) to the client. Tool calls, like database queries often contain large amounts of data or sensitive information that should not be visible on the client. This library provides a type-safe filter to apply selective streaming of certain message parts.

Installation

This library only supports AI SDK v5.

npm install ai-filter-stream

Usage

Use the filterUIMessageStream function to wrap the UI message stream from result.toUIMessageStream() and provide a filter to include or exclude certain UI message parts:

[!NOTE]
Providing a MyUIMessage type to filterUIMessageStream<MyMessage>() is optional and only required for type-safety so that the part type is inferred based on your tools and data parts.

import { streamText } from 'ai';
import { filterUIMessageStream } from 'ai-filter-stream';
import type { UIMessage, InferUITools } from 'ai';

type MyUIMessageMetadata = {};

type MyDataPart = {};

type MyTools = InferUITools<typeof tools>;

// Define your UI message type for type safety
// See: https://ai-sdk.dev/docs/reference/ai-sdk-core/ui-message
type MyUIMessage = UIMessage<
  MyUIMessageMetadata, // or unknown
  MyDataPart, // or unknown
  MyTools,
>;

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,
    }),
  }),
};

const result = streamText({
  model,
  prompt: 'What is the weather in Tokyo?',
  tools,
});

// Inclusive filtering: include only `text` parts
const stream = filterUIMessageStream<MyMessage>(result.toUIMessageStream(), {
  includeParts: ['text'], // Autocomplete works here!
});

// Exclusive filtering: exclude only `tool-weather` parts
const stream = filterUIMessageStream<MyMessage>(result.toUIMessageStream(), {
  excludeParts: ['reasoning', 'tool-calculator'], // Autocomplete works here!
});

// Dynamic filtering: apply filter function for each chunk
const stream = filterUIMessageStream<MyMessage>(result.toUIMessageStream(), {
  filterParts: ({ partType }) => {
    // Always include text
    if (partType === 'text') return true;

    // Only include tools that start with 'weather'
    if (partType.startsWith('tool-weather')) return true;

    // Exclude everything else
    return false;
  },
});

Part Type Mapping

The filter operates on UIMessagePart part types, which are derived from UIMessageChunk chunk types:

| Part Type | Chunk Types | | ----------------- | ------------------------------------- | | text | text-start, text-delta, text-end | | reasoning | reasoning-start, reasoning-delta, reasoning-end | | tool-{name} | tool-input-start, tool-input-delta, tool-input-available, tool-input-error, tool-output-available, tool-output-error | | data-{name} | data-{name} | | step-start | start-step | | file | file | | source-url | source-url | | source-document | source-document |

Controls chunks are always passed through regardless of filter settings:

• start: Stream start marker
• finish: Stream finish marker
• abort: Stream abort marker
• message-metadata: Message metadata updates
• error: Error messages

Start-Step Filtering

The filter automatically handles step boundaries, that means a start-step is only emitted if the actual content is not filtered:

1. start-step is buffered until the first content chunk is encountered
2. If the first content chunk passes the filter, start-step is included
3. If the first content chunk is filtered out, start-step is also filtered out
4. finish-step is only included if the corresponding start-step was included

Type Safety

The toUIMessageStream() from streamText() returns a generic stream ReadableStream<UIMessageChunk> which means that the original UIMessage cannot be inferred automatically. To enable autocomplete and type-safety for filtering parts by type, we need to pass our own UIMessage as generic param to filterUIMessageStream():

type MyMessage = UIMessage<MyMetadata, MyData, MyTools>;

const stream = filterUIMessageStream<MyMessage>(
  result.toUIMessageStream(), // returns generic ReadableStream<UIMessageChunk>
  {
    includeParts: ['text', 'tool-weather'] }, // type-safe through MyMessage
  }
);

See UIMessage in the AI SDK docs to create your own UIMessage type.

API Reference

filterUIMessageStream

function filterUIMessageStream<UI_MESSAGE extends UIMessage>(
  stream: ReadableStream<UIMessageChunk>,
  options: FilterUIMessageStreamOptions<UI_MESSAGE>,
): AsyncIterableStream<InferUIMessageChunk<UI_MESSAGE>>

FilterUIMessageStreamOptions

type FilterUIMessageStreamOptions<UI_MESSAGE extends UIMessage> =
  | {
      filterParts: (options: { partType: InferUIMessagePartType<UI_MESSAGE> }) => boolean;
    }
  | {
      includeParts: Array<InferUIMessagePartType<UI_MESSAGE>>;
    }
  | {
      excludeParts: Array<InferUIMessagePartType<UI_MESSAGE>>;
    };