vercel / ai

Build AI-powered applications with React, Svelte, Vue, and Solid
https://sdk.vercel.ai/docs
Other
10.11k stars 1.51k forks source link

useChat hook's returned helper `messages` returned as new empty array every time #615

Closed davit-b closed 1 year ago

davit-b commented 1 year ago

Description

Use the code example from the AI SDK for React using useChat. The problem is simple:

  const {
    messages
  } = useChat({...})

This hook instantiates an empty array for messages=[] every time the component where useChat is use is re-render.

Now imagine a 2-level component structure

<ParentComponent /> (useChat)
      <FirstChild  messages={messages} />
      <SecondChild /> 

If the SecondChild component's triggers a re-render of the ParentComponent for any reason, the messages object is re-instantiated as an empty array and this triggers the FirstChild component to re-render.

EVEN if you wrap FirstChild with React.memo(...) where the props in the 1st call are messages=[] and the 2nd call are messages=[], it will trigger a render because the shallow comparison of props shows it as a new empty array.

The workaround (for anyone else running into this)

  const [guardedMessages, setGuardedMessages] = useState<...>([])

  useEffect(() => {
    if (!Array.isArray(messages) || !messages.length) return
    setGuardedMessages(messages)
  }, [messages])

// with a structure like
  <ParentComponent /> (useChat)
      <FirstChild  messages={guardedMessages} />
      <SecondChild /> 

After implementing the workaround, I don't get the to re-render a million times whenever causes to trigger a re-render because of a small state update.

I really believe the SDK shouldn't return a new empty array each instantiation. Yes, as users of the SDK we can contort our parent-child hierarchy into a specific way to make it work.... but I really think when messages=[] it should return the same empty array rather than re-creating it over and over.

Code at fault is: https://github.com/vercel/ai/blame/88451356c25b219a37946fb3b75da4f0d936a72c/packages/core/react/use-chat.ts#L623

Thoughts?

Code example

No response

Additional context

No response

davit-b commented 1 year ago

I believe I ran into this issue because I see the re-render happening when messages is empty. This is the case because I only use messages as a buffer of messages to make a request to api/chat, after the response streams I clear the buffer to a larger state that holds all my chat messages.

This way, I avoid using messages as my primary state and every network call I do not need to send the ENTIRE message block to /api/chat, only the user's recent input

davit-b commented 1 year ago

Thinking about it, I guess I could implement a custom check in my memo use to check if the oldProps.messages === [] and the newProps.messages === [], rather than the shallow check which I believe checks if the reference is the same and two empty arrays don't have the same reference.

https://react.dev/reference/react/memo#specifying-a-custom-comparison-function

MaxLeiter commented 1 year ago

A PR fixing messages stability would be great.

davit-b commented 1 year ago

I'm not sure how. I'm not even sure of the root cause yet: if it's the hook being reinstantiated or it's related the to useEffect inside of useChat. Honestly, glancing through this sdk for a few months now has been humbling - there's a lot about JS and React I haven't learned yet.

I'd have to pull the package and investigate deeper.

lgrammel commented 1 year ago

You can use memo with a custom comparison function as follows:

'use client';

import { Message, useChat } from 'ai/react';
import { memo, useRef } from 'react';

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

  return (
    <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
      <MessageView messages={messages} />

      <form onSubmit={handleSubmit}>
        <input
          className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
          value={input}
          placeholder="Say something..."
          onChange={handleInputChange}
        />
      </form>
    </div>
  );
}

const MessageView = memo(
  function ({ messages }: { messages: Message[] }) {
    const renderCounter = useRef(0);
    renderCounter.current = renderCounter.current + 1;

    return (
      <>
        <h1>Renders: {renderCounter.current}</h1>
        {messages.map(m => (
          <div key={m.id} className="whitespace-pre-wrap">
            {m.role === 'user' ? 'User: ' : 'AI: '}
            {m.content}
          </div>
        ))}
      </>
    );
  },
   // custom function that checks if both message [] are of length 0:
  (prev, next) =>
    (prev.messages.length === 0 && next.messages.length === 0) ||
    prev.messages === next.messages,
);
lgrammel commented 1 year ago

I've traced the changes of the messages array. The root of the changing arrays is in the function call defaults: https://github.com/vercel/ai/blame/88451356c25b219a37946fb3b75da4f0d936a72c/packages/core/react/use-chat.ts#L339

The || [] in the return statement was not executed in my traces, bc messages is always defined at that point.

davit-b commented 1 year ago

@lgrammel thank you!!