How to memoize messages properly when using "useChat"

[LaurinBrechter]

When I am using useChat and then rendering the messages in the ui as in the example below, all messages get rerendered whenever a single chunk is being streamed back to the client. This can be very costly when rendering the messages is expensive. Are there any examples of how to properly memoize 'old' messages? I have something working where I put past messages into a different state that doesn't get rerendered on every chunk from the response but I was wondering if anyone else was having this problem when simply rendering the messages array

'use client';

import { useChat } from 'ai/react';

export default function Page() {
  const { messages, input, handleInputChange, handleSubmit } = useChat({});

  return (
    <>
      // rerenders on every chunk
      {messages.map(message => (
        <div key={message.id}>
          {message.role === 'user' ? 'User: ' : 'AI: '}
          {message.content}
        </div>
      ))}

      <form onSubmit={handleSubmit}>
        <input name="prompt" value={input} onChange={handleInputChange} />
        <button type="submit">Submit</button>
      </form>
    </>
  );
}

[C-W-D-Harshit]

Hi did you find any solution for this?

[LaurinBrechter]

Hey, so from my understanding since the messages returned by useChat is state you can't really memoize it since it changes with each chunk. I therefore create a second state that stores all completed messages. These old messages only get updated once the generation is finished (at least better than at each chunk). I then create two ChatList components that I render (all client-side). One is for the old messages and is memoized. The other is for the new messages. Since the new messages also contain the old ones that I am already rendering I do the slice to exclude the old messages. There's probably a better way to do this but I am not really a memoization expert.

// 
export default function ChatWindow(props: {
  readOnly: boolean;
  threadId: string;
  locked: boolean;
  initialMessages: Array<Message>;
  chatType: string;
}) {
  const [oldMessages, setOldMessages] = useState(props.initialMessages);
  const { messages, handleSubmit, input, setInput, append, isLoading, stop, error } = useChat(chatConfig);
  useEffect(() => {
      if (!isLoading) {
        setOldMessages([...messages]);
      }
    }, [isLoading, messages]);
  const memoizedOldChatList = useMemo(() => (
      <ChatList
        messages={oldMessages}
        isShared={false}
        isLoading={isLoading}
        memoized={true}
      />
    ), [oldMessages, isLoading]);
  
    const memoizedNewChatList = useMemo(() => (
      <ChatList

        // remove the old messages that we are already rendering
        messages={messages.slice(oldMessages.length)}
        isShared={false}
        isLoading={isLoading}
        memoized={false}
      />
    ), [messages, oldMessages, isLoading]);

  // ... render both ChatList components

[EducatedAI-kaisera]

Check the Vercel ChatBot template and here: https://sdk.vercel.ai/cookbook/next/markdown-chatbot-with-memoization#memoized-markdown-component

[lgrammel]

Check out our recipe: https://sdk.vercel.ai/cookbook/next/markdown-chatbot-with-memoization