Implement custom id generator parameter (#733)

Co-authored-by: Max Leiter <max.leiter@vercel.com>

.changeset/eighty-spies-grow.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+Add `generateId` to use-chat params to allow overriding message ID generation

packages/core/react/use-chat.ts

@@ -9,6 +9,7 @@ import type {
   JSONValue,
   Message,
   UseChatOptions,
+  IdGenerator,
 } from '../shared/types';
 
 import { callApi } from '../shared/call-api';
@@ -88,6 +89,7 @@ const getStreamedResponse = async (
   extraMetadataRef: React.MutableRefObject<any>,
   messagesRef: React.MutableRefObject<Message[]>,
   abortControllerRef: React.MutableRefObject<AbortController | null>,
+  generateId: IdGenerator,
   onFinish?: (message: Message) => void,
   onResponse?: (response: Response) => void | Promise<void>,
   sendExtraMessageFields?: boolean,
@@ -111,7 +113,7 @@ const getStreamedResponse = async (
   if (typeof api !== 'string') {
     // In this case, we are handling a Server Action. No complex mode handling needed.
 
-    const replyId = nanoid();
+    const replyId = generateId();
     const createdAt = new Date();
     let responseMessage: Message = {
       id: replyId,
@@ -185,6 +187,7 @@ const getStreamedResponse = async (
       mutateStreamData([...(existingData || []), ...(data || [])], false);
     },
     onFinish,
+    generateId,
   });
 };
 
@@ -201,6 +204,7 @@ export function useChat({
   credentials,
   headers,
   body,
+  generateId = nanoid,
 }: Omit<UseChatOptions, 'api'> & {
   api?: string | StreamingReactResponseAction;
 } = {}): UseChatHelpers {
@@ -274,6 +278,7 @@ export function useChat({
               extraMetadataRef,
               messagesRef,
               abortControllerRef,
+              generateId,
               onFinish,
               onResponse,
               sendExtraMessageFields,
@@ -317,6 +322,7 @@ export function useChat({
       experimental_onFunctionCall,
       messagesRef.current,
       abortControllerRef.current,
+      generateId,
     ],
   );
 
@@ -326,7 +332,7 @@ export function useChat({
       { options, functions, function_call, data }: ChatRequestOptions = {},
     ) => {
       if (!message.id) {
-        message.id = nanoid();
+        message.id = generateId();
       }
 
       const chatRequest: ChatRequest = {
@@ -339,7 +345,7 @@ export function useChat({
 
       return triggerRequest(chatRequest);
     },
-    [triggerRequest],
+    [triggerRequest, generateId],
   );
 
   const reload = useCallback(

packages/core/shared/call-api.ts

@@ -1,6 +1,6 @@
 import { nanoid } from 'nanoid';
 import { parseComplexResponse } from './parse-complex-response';
-import { FunctionCall, JSONValue, Message } from './types';
+import { FunctionCall, IdGenerator, JSONValue, Message } from './types';
 import { COMPLEX_HEADER, createChunkDecoder } from './utils';
 
 export async function callApi({
@@ -15,6 +15,7 @@ export async function callApi({
   onResponse,
   onUpdate,
   onFinish,
+  generateId,
 }: {
   api: string;
   messages: Omit<Message, 'id'>[];
@@ -27,6 +28,7 @@ export async function callApi({
   onResponse?: (response: Response) => void | Promise<void>;
   onUpdate: (merged: Message[], data: JSONValue[] | undefined) => void;
   onFinish?: (message: Message) => void;
+  generateId: IdGenerator;
 }) {
   const response = await fetch(api, {
     method: 'POST',
@@ -78,14 +80,15 @@ export async function callApi({
           onFinish(prefixMap.text);
         }
       },
+      generateId,
     });
   } else {
     const createdAt = new Date();
     const decode = createChunkDecoder(false);
 
     // TODO-STREAMDATA: Remove this once Stream Data is not experimental
     let streamedResponse = '';
-    const replyId = nanoid();
+    const replyId = generateId();
     let responseMessage: Message = {
       id: replyId,
       createdAt,

packages/core/shared/types.ts

@@ -40,6 +40,8 @@ interface Function {
   description?: string;
 }
 
+export type IdGenerator = () => string;
+
 /**
  * Shared types between the API and UI packages.
  */
@@ -137,6 +139,12 @@ export type UseChatOptions = {
    */
   onError?: (error: Error) => 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.
+   */
+  generateId?: IdGenerator;
+
   /**
    * The credentials mode to be used for the fetch request.
    * Possible values are: 'omit', 'same-origin', 'include'.

packages/core/solid/use-chat.ts

@@ -80,6 +80,7 @@ export function useChat({
   credentials,
   headers,
   body,
+  generateId = nanoid,
 }: UseChatOptions = {}): UseChatHelpers {
   // Generate a unique ID for the chat if not provided.
   const chatId = id || `chat-${uniqueId++}`;
@@ -174,6 +175,7 @@ export function useChat({
                 mutate(previousMessages.data);
               }
             },
+            generateId,
           });
         },
         experimental_onFunctionCall,
@@ -203,7 +205,7 @@ export function useChat({
 
   const append: UseChatHelpers['append'] = async (message, options) => {
     if (!message.id) {
-      message.id = nanoid();
+      message.id = generateId();
     }
     return triggerRequest(
       (messages() ?? []).concat(message as Message),

packages/core/streams/streaming-react-response.ts

@@ -9,8 +9,8 @@
  */
 
 import { parseComplexResponse } from '../shared/parse-complex-response';
-import { JSONValue } from '../shared/types';
-import { createChunkDecoder } from '../shared/utils';
+import { IdGenerator, JSONValue } from '../shared/types';
+import { createChunkDecoder, nanoid } from '../shared/utils';
 import { experimental_StreamData } from './stream-data';
 
 type UINode = string | JSX.Element | JSX.Element[] | null | undefined;
@@ -36,6 +36,7 @@ export class experimental_StreamingReactResponse {
         data?: JSONValue[] | undefined;
       }) => UINode | Promise<UINode>;
       data?: experimental_StreamData;
+      generateId?: IdGenerator;
     },
   ) {
     let resolveFunc: (row: ReactResponseRow) => void = () => {};
@@ -70,6 +71,7 @@ export class experimental_StreamingReactResponse {
 
           lastPayload = payload;
         },
+        generateId: options.generateId ?? nanoid,
         onFinish: () => {
           // The last payload is resolved twice. This is necessary because we immediately
           // push out a payload, but we also need to forward the finish event with a payload.

packages/core/svelte/use-chat.ts

@@ -9,6 +9,7 @@ import type {
   JSONValue,
   Message,
   UseChatOptions,
+  IdGenerator,
 } from '../shared/types';
 import { nanoid } from '../shared/utils';
 export type { CreateMessage, Message, UseChatOptions };
@@ -70,6 +71,7 @@ const getStreamedResponse = async (
   },
   previousMessages: Message[],
   abortControllerRef: AbortController | null,
+  generateId: IdGenerator,
   onFinish?: (message: Message) => void,
   onResponse?: (response: Response) => void | Promise<void>,
   sendExtraMessageFields?: boolean,
@@ -120,6 +122,7 @@ const getStreamedResponse = async (
       mutateStreamData([...(existingData || []), ...(data || [])]);
     },
     onFinish,
+    generateId,
   });
 };
 
@@ -140,6 +143,7 @@ export function useChat({
   credentials,
   headers,
   body,
+  generateId = nanoid,
 }: UseChatOptions = {}): UseChatHelpers {
   // Generate a unique id for the chat if not provided.
   const chatId = id || `chat-${uniqueId++}`;
@@ -201,6 +205,7 @@ export function useChat({
             extraMetadata,
             get(messages),
             abortController,
+            generateId,
             onFinish,
             onResponse,
             sendExtraMessageFields,
@@ -237,7 +242,7 @@ export function useChat({
     { options, functions, function_call }: ChatRequestOptions = {},
   ) => {
     if (!message.id) {
-      message.id = nanoid();
+      message.id = generateId();
     }
 
     const chatRequest: ChatRequest = {

packages/core/vue/use-chat.ts

@@ -74,6 +74,7 @@ export function useChat({
   credentials,
   headers,
   body,
+  generateId = nanoid,
 }: UseChatOptions = {}): UseChatHelpers {
   // Generate a unique ID for the chat if not provided.
   const chatId = id || `chat-${uniqueId++}`;
@@ -173,6 +174,7 @@ export function useChat({
               // Restore the previous messages if the request fails.
               mutate(previousMessages);
             },
+            generateId,
           });
         },
         experimental_onFunctionCall,
@@ -202,7 +204,7 @@ export function useChat({
 
   const append: UseChatHelpers['append'] = async (message, options) => {
     if (!message.id) {
-      message.id = nanoid();
+      message.id = generateId();
     }
     return triggerRequest(messages.value.concat(message as Message), options);
   };