feat add onParts support for chat API and processing functions

packages/react/src/use-chat.ts

@@ -127,6 +127,7 @@ export function useChat({
   streamProtocol = 'data',
   onResponse,
   onFinish,
+  onParts,
   onError,
   credentials,
   headers,
@@ -152,7 +153,7 @@ export function useChat({
     messages: UIMessage[];
     requestData?: JSONValue;
     requestBody?: object;
-  }) => unknown;
+  }) => Promise<unknown> | unknown;
 
   /**
 Custom throttle wait in ms for the chat messages and data updates.
@@ -298,12 +299,12 @@ By default, it's set to 1, which means that only a single LLM call is made.
 
         await callChatApi({
           api,
-          body: experimental_prepareRequestBody?.({
+          body: (await experimental_prepareRequestBody?.({
             id: chatId,
             messages: chatMessages,
             requestData: chatRequest.data,
             requestBody: chatRequest.body,
-          }) ?? {
+          })) ?? {
             id: chatId,
             messages: constructedMessagesPayload,
             data: chatRequest.data,
@@ -322,6 +323,7 @@ By default, it's set to 1, which means that only a single LLM call is made.
               throttledMutate(previousMessages, false);
             }
           },
+          onParts,
           onResponse,
           onUpdate({ message, data, replaceLastMessage }) {
             mutateStatus('streaming');
@@ -390,6 +392,7 @@ By default, it's set to 1, which means that only a single LLM call is made.
       api,
       extraMetadataRef,
       onResponse,
+      onParts,
       onFinish,
       onError,
       setError,
@@ -418,11 +421,12 @@ By default, it's set to 1, which means that only a single LLM call is made.
         headers,
         body,
         experimental_attachments = message.experimental_attachments,
+        experimental_attachmentGenerator,
       }: ChatRequestOptions = {},
     ) => {
-      const attachmentsForRequest = await prepareAttachmentsForRequest(
-        experimental_attachments,
-      );
+      const attachmentsForRequest = experimental_attachmentGenerator
+        ? await experimental_attachmentGenerator()
+        : await prepareAttachmentsForRequest(experimental_attachments);
 
       const messages = messagesRef.current.concat({
         ...message,
@@ -522,9 +526,9 @@ By default, it's set to 1, which means that only a single LLM call is made.
         };
       }
 
-      const attachmentsForRequest = await prepareAttachmentsForRequest(
-        options.experimental_attachments,
-      );
+      const attachmentsForRequest = options?.experimental_attachmentGenerator
+        ? await options.experimental_attachmentGenerator()
+        : await prepareAttachmentsForRequest(options.experimental_attachments);
 
       const messages = messagesRef.current.concat({
         id: generateId(),

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

@@ -1,6 +1,12 @@
 import { processChatResponse } from './process-chat-response';
 import { processChatTextResponse } from './process-chat-text-response';
-import { IdGenerator, JSONValue, UIMessage, UseChatOptions } from './types';
+import {
+  IdGenerator,
+  JSONValue,
+  onParts,
+  UIMessage,
+  UseChatOptions,
+} from './types';
 
 // use function to allow for mocking in tests:
 const getOriginalFetch = () => fetch;
@@ -13,6 +19,7 @@ export async function callChatApi({
   headers,
   abortController,
   restoreMessagesOnFailure,
+  onParts,
   onResponse,
   onUpdate,
   onFinish,
@@ -29,6 +36,7 @@ export async function callChatApi({
   headers: HeadersInit | undefined;
   abortController: (() => AbortController | null) | undefined;
   restoreMessagesOnFailure: () => void;
+  onParts?: onParts;
   onResponse: ((response: Response) => void | Promise<void>) | undefined;
   onUpdate: (options: {
     message: UIMessage;
@@ -105,9 +113,10 @@ export async function callChatApi({
         update: onUpdate,
         lastMessage,
         onToolCall,
-        onFinish({ message, finishReason, usage }) {
+        onParts,
+        async onFinish({ message, finishReason, usage }) {
           if (onFinish && message != null) {
-            onFinish(message, { usage, finishReason });
+            await onFinish(message, { usage, finishReason });
           }
         },
         generateId,

packages/ui-utils/src/process-chat-response.ts

@@ -7,20 +7,23 @@ import {
 import { parsePartialJson } from './parse-partial-json';
 import { processDataStream } from './process-data-stream';
 import type {
+  FileUIPart,
   JSONValue,
   ReasoningUIPart,
   TextUIPart,
   ToolInvocation,
   ToolInvocationUIPart,
   UIMessage,
   UseChatOptions,
+  onParts,
 } from './types';
 
 export async function processChatResponse({
   stream,
   update,
   onToolCall,
   onFinish,
+  onParts,
   generateId = generateIdFunction,
   getCurrentDate = () => new Date(),
   lastMessage,
@@ -36,7 +39,8 @@ export async function processChatResponse({
     message: UIMessage | undefined;
     finishReason: LanguageModelV1FinishReason;
     usage: LanguageModelUsage;
-  }) => void;
+  }) => void | Promise<void>;
+  onParts?: onParts;
   generateId?: () => string;
   getCurrentDate?: () => Date;
   lastMessage: UIMessage | undefined;
@@ -200,12 +204,19 @@ export async function processChatResponse({
 
       execUpdate();
     },
-    onFilePart(value) {
-      message.parts.push({
-        type: 'file',
-        mimeType: value.mimeType,
-        data: value.data,
-      });
+    async onFilePart(value) {
+      const part: FileUIPart = onParts?.onFilePart
+        ? await onParts?.onFilePart({
+            mimeType: value.mimeType,
+            data: value.data,
+          })
+        : {
+            type: 'file',
+            mimeType: value.mimeType,
+            data: value.data,
+          };
+
+      message.parts.push(part);
 
       execUpdate();
     },
@@ -384,5 +395,5 @@ export async function processChatResponse({
     },
   });
 
-  onFinish?.({ message, finishReason, usage });
+  await onFinish?.({ message, finishReason, usage });
 }

packages/ui-utils/src/process-chat-text-response.ts

@@ -46,7 +46,7 @@ export async function processChatTextResponse({
   });
 
   // in text mode, we don't have usage information or finish reason:
-  onFinish?.(resultMessage, {
+  await onFinish?.(resultMessage, {
     usage: { completionTokens: NaN, promptTokens: NaN, totalTokens: NaN },
     finishReason: 'unknown',
   });

packages/ui-utils/src/types.ts

@@ -4,6 +4,7 @@ import {
 } from '@ai-sdk/provider';
 import { FetchFunction, ToolCall, ToolResult } from '@ai-sdk/provider-utils';
 import { LanguageModelUsage } from './duplicated/usage';
+import { DataStreamPartType } from './data-stream-parts';
 
 export * from './use-assistant-types';
 
@@ -268,6 +269,8 @@ Additional data to be sent to the API endpoint.
    * Allow submitting an empty message. Defaults to `false`.
    */
   allowEmptySubmit?: boolean;
+
+  experimental_attachmentGenerator?: () => Promise<Attachment[]> | Attachment[];
 };
 
 export type UseChatOptions = {
@@ -332,7 +335,7 @@ either synchronously or asynchronously.
       usage: LanguageModelUsage;
       finishReason: LanguageModelV1FinishReason;
     },
-  ) => void;
+  ) => void | Promise<void>;
 
   /**
    * Callback function to be called when an error is encountered.
@@ -388,6 +391,11 @@ Custom fetch implementation. You can use it as a middleware to intercept request
 or to provide a custom fetch implementation for e.g. testing.
     */
   fetch?: FetchFunction;
+
+  /**
+   * Callback function to be called when a part is received while streaming.
+   */
+  onParts?: onParts;
 };
 
 export type UseCompletionOptions = {
@@ -501,3 +509,9 @@ export type DataMessage = {
   role: 'data';
   data: JSONValue; // application-specific data
 };
+
+export type onParts = {
+  onFilePart?: (
+    streamPart: (DataStreamPartType & { type: 'file' })['value'],
+  ) => Promise<FileUIPart> | FileUIPart;
+};