feat (ai/ui): add mutator function support to useChat / setMessages. (#…

…2347)

.changeset/tough-spies-buy.md

@@ -0,0 +1,9 @@
+---
+'@ai-sdk/svelte': patch
+'@ai-sdk/react': patch
+'@ai-sdk/solid': patch
+'ai': patch
+'@ai-sdk/vue': patch
+---
+
+feat (ai/ui): add mutator function support to useChat / setMessages

content/docs/07-reference/ai-sdk-ui/01-use-chat.mdx

@@ -388,7 +388,7 @@ Allows you to easily create a conversational user interface for your chatbot app
     },
     {
       name: 'setMessages',
-      type: '(messages: Message[]) => void',
+      type: '(messages: Message[] | ((messages: Message[]) => Message[]) => void',
       description:
         'Function to update the `messages` state locally without triggering an API call.',
     },

packages/core/svelte/use-chat.ts

@@ -49,7 +49,10 @@ export type UseChatHelpers = {
    * edit the messages on the client, and then trigger the `reload` method
    * manually to regenerate the AI response.
    */
-  setMessages: (messages: Message[]) => void;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+
   /** The current value of the input */
   input: Writable<string>;
   /** Form submission handler to automatically reset input and append a user message  */
@@ -359,8 +362,14 @@ export function useChat({
     }
   };
 
-  const setMessages = (messages: Message[]) => {
-    mutate(messages);
+  const setMessages = (
+    messagesArg: Message[] | ((messages: Message[]) => Message[]),
+  ) => {
+    if (typeof messagesArg === 'function') {
+      messagesArg = messagesArg(get(messages));
+    }
+
+    mutate(messagesArg);
   };
 
   const input = writable(initialInput);

packages/react/src/use-chat.ts

@@ -51,7 +51,9 @@ export type UseChatHelpers = {
    * edit the messages on the client, and then trigger the `reload` method
    * manually to regenerate the AI response.
    */
-  setMessages: (messages: Message[]) => void;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
   /** The current value of the input */
   input: string;
   /** setState-powered method to update the input value */
@@ -507,7 +509,11 @@ By default, it's set to 0, which will disable the feature.
   }, []);
 
   const setMessages = useCallback(
-    (messages: Message[]) => {
+    (messages: Message[] | ((messages: Message[]) => Message[])) => {
+      if (typeof messages === 'function') {
+        messages = messages(messagesRef.current);
+      }
+
       mutate(messages, false);
       messagesRef.current = messages;
     },

packages/solid/src/use-chat.ts

@@ -58,7 +58,9 @@ export type UseChatHelpers = {
    * edit the messages on the client, and then trigger the `reload` method
    * manually to regenerate the AI response.
    */
-  setMessages: (messages: Message[]) => void;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
   /** The current value of the input */
   input: Accessor<string>;
   /** Signal setter to update the input value */
@@ -362,9 +364,15 @@ export function useChat(
     }
   };
 
-  const setMessages = (messages: Message[]) => {
-    mutate(messages);
-    messagesRef = messages;
+  const setMessages = (
+    messagesArg: Message[] | ((messages: Message[]) => Message[]),
+  ) => {
+    if (typeof messagesArg === 'function') {
+      messagesArg = messagesArg(messagesRef);
+    }
+
+    mutate(messagesArg);
+    messagesRef = messagesArg;
   };
 
   const [input, setInput] = createSignal(

packages/svelte/src/use-chat.ts

@@ -49,7 +49,10 @@ export type UseChatHelpers = {
    * edit the messages on the client, and then trigger the `reload` method
    * manually to regenerate the AI response.
    */
-  setMessages: (messages: Message[]) => void;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+
   /** The current value of the input */
   input: Writable<string>;
   /** Form submission handler to automatically reset input and append a user message  */
@@ -356,8 +359,14 @@ export function useChat({
     }
   };
 
-  const setMessages = (messages: Message[]) => {
-    mutate(messages);
+  const setMessages = (
+    messagesArg: Message[] | ((messages: Message[]) => Message[]),
+  ) => {
+    if (typeof messagesArg === 'function') {
+      messagesArg = messagesArg(get(messages));
+    }
+
+    mutate(messagesArg);
   };
 
   const input = writable(initialInput);

packages/vue/src/use-chat.ts

@@ -47,7 +47,9 @@ export type UseChatHelpers = {
    * edit the messages on the client, and then trigger the `reload` method
    * manually to regenerate the AI response.
    */
-  setMessages: (messages: Message[]) => void;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
   /** The current value of the input */
   input: Ref<string>;
   /** Form submission handler to automatically reset input and append a user message  */
@@ -264,8 +266,14 @@ export function useChat({
     }
   };
 
-  const setMessages = (messages: Message[]) => {
-    mutate(messages);
+  const setMessages = (
+    messagesArg: Message[] | ((messages: Message[]) => Message[]),
+  ) => {
+    if (typeof messagesArg === 'function') {
+      messagesArg = messagesArg(messages.value);
+    }
+
+    mutate(messagesArg);
   };
 
   const input = ref(initialInput);