useChat chunk

packages/react/src/use-chat.ts

@@ -87,6 +87,7 @@ const getStreamedResponse = async (
   abortControllerRef: React.MutableRefObject<AbortController | null>,
   generateId: IdGenerator,
   streamProtocol: UseChatOptions['streamProtocol'],
+  onChunk: UseChatOptions['onChunk'],
   onFinish: UseChatOptions['onFinish'],
   onResponse: ((response: Response) => void | Promise<void>) | undefined,
   onToolCall: UseChatOptions['onToolCall'] | undefined,
@@ -178,6 +179,7 @@ const getStreamedResponse = async (
       mutateStreamData([...(existingData || []), ...(data || [])], false);
     },
     onToolCall,
+    onChunk,
     onFinish,
     generateId,
     fetch,
@@ -203,6 +205,7 @@ export function useChat({
   onResponse,
   onFinish,
   onError,
+  onChunk,
   credentials,
   headers,
   body,
@@ -363,6 +366,7 @@ By default, it's set to 1, which means that only a single LLM call is made.
               abortControllerRef,
               generateId,
               streamProtocol,
+              onChunk,
               onFinish,
               onResponse,
               onToolCall,
@@ -420,6 +424,7 @@ By default, it's set to 1, which means that only a single LLM call is made.
       api,
       extraMetadataRef,
       onResponse,
+      onChunk,
       onFinish,
       onError,
       setError,

packages/react/src/use-chat.ui.test.tsx

@@ -1,6 +1,7 @@
 /* eslint-disable @next/next/no-img-element */
 import { withTestServer } from '@ai-sdk/provider-utils/test';
 import {
+  StreamPartType,
   formatStreamPart,
   getTextFromDataUrl,
   Message,
@@ -30,13 +31,18 @@ describe('stream data stream', () => {
     };
   }> = [];
 
+  let chunks: StreamPartType[] = [];
+
   const TestComponent = () => {
     const [id, setId] = React.useState<string>('first-id');
     const { messages, append, error, data, isLoading } = useChat({
       id,
       onFinish: (message, options) => {
         onFinishCalls.push({ message, options });
       },
+      onChunk(event) {
+        chunks.push(event.chunk);
+      },
     });
 
     return (
@@ -70,12 +76,14 @@ describe('stream data stream', () => {
   beforeEach(() => {
     render(<TestComponent />);
     onFinishCalls = [];
+    chunks = [];
   });
 
   afterEach(() => {
     vi.restoreAllMocks();
     cleanup();
     onFinishCalls = [];
+    chunks = [];
   });
 
   it(
@@ -257,6 +265,46 @@ describe('stream data stream', () => {
       ),
     );
   });
+
+  it(
+    'options.onChunk',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: [
+          formatStreamPart('text', 'Hello'),
+          formatStreamPart('text', ','),
+          formatStreamPart('text', ' world'),
+          formatStreamPart('text', '.'),
+        ],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append'));
+
+        await screen.findByTestId('message-1');
+
+        expect(chunks).toStrictEqual([
+          {
+            type: 'text',
+            value: 'Hello',
+          },
+          {
+            type: 'text',
+            value: ',',
+          },
+          {
+            type: 'text',
+            value: ' world',
+          },
+          {
+            type: 'text',
+            value: '.',
+          },
+        ]);
+      },
+    ),
+  );
 });
 
 describe('text stream', () => {

packages/ui-utils/src/call-chat-api.ts

@@ -15,6 +15,7 @@ export async function callChatApi({
   restoreMessagesOnFailure,
   onResponse,
   onUpdate,
+  onChunk,
   onFinish,
   onToolCall,
   generateId,
@@ -29,6 +30,7 @@ export async function callChatApi({
   restoreMessagesOnFailure: () => void;
   onResponse: ((response: Response) => void | Promise<void>) | undefined;
   onUpdate: (newMessages: Message[], data: JSONValue[] | undefined) => void;
+  onChunk?: UseChatOptions['onChunk'];
   onFinish: UseChatOptions['onFinish'];
   onToolCall: UseChatOptions['onToolCall'];
   generateId: IdGenerator;
@@ -116,6 +118,7 @@ export async function callChatApi({
         abortControllerRef:
           abortController != null ? { current: abortController() } : undefined,
         update: onUpdate,
+        onChunk,
         onToolCall,
         onFinish({ message, finishReason, usage }) {
           if (onFinish && message != null) {

packages/ui-utils/src/index.ts

@@ -18,4 +18,4 @@ export { readDataStream } from './read-data-stream';
 export { asSchema, jsonSchema, zodSchema } from './schema';
 export type { Schema } from './schema';
 export { formatStreamPart, parseStreamPart } from './stream-parts';
-export type { StreamPart, StreamString } from './stream-parts';
+export type { StreamPart, StreamString, StreamPartType } from './stream-parts';

packages/ui-utils/src/process-data-protocol-response.ts

@@ -36,6 +36,7 @@ export async function processDataProtocolResponse({
   reader,
   abortControllerRef,
   update,
+  onChunk,
   onToolCall,
   onFinish,
   generateId = generateIdFunction,
@@ -47,6 +48,7 @@ export async function processDataProtocolResponse({
   };
   update: (newMessages: Message[], data: JSONValue[] | undefined) => void;
   onToolCall?: UseChatOptions['onToolCall'];
+  onChunk?: UseChatOptions['onChunk'];
   onFinish?: (options: {
     message: Message | undefined;
     finishReason: LanguageModelV1FinishReason;
@@ -89,9 +91,12 @@ export async function processDataProtocolResponse({
   let finishReason: LanguageModelV1FinishReason = 'unknown';
 
   // we create a map of each prefix, and for each prefixed message we push to the map
-  for await (const { type, value } of readDataStream(reader, {
+  for await (const chunk of readDataStream(reader, {
     isAborted: () => abortControllerRef?.current === null,
   })) {
+    const { type, value } = chunk;
+    await onChunk?.({ chunk });
+
     if (type === 'error') {
       throw new Error(value);
     }

packages/ui-utils/src/types.ts

@@ -3,6 +3,7 @@ import { FetchFunction } from '@ai-sdk/provider-utils';
 import { CompletionTokenUsage } from './duplicated/token-usage';
 import { ToolCall as CoreToolCall } from './duplicated/tool-call';
 import { ToolResult as CoreToolResult } from './duplicated/tool-result';
+import { StreamPartType } from './stream-parts';
 
 export * from './use-assistant-types';
 
@@ -448,6 +449,11 @@ either synchronously or asynchronously.
    */
   onError?: (error: Error) => void;
 
+  /**
+   * Callback that is called for each chunk of the stream.
+   */
+  onChunk?: (event: { chunk: StreamPartType }) => Promise<void> | void;
+
   /**
    * A way to provide a function that is going to be used for ids for messages.
    * If not provided nanoid is used by default.