refactor: replace isLoading with status (#861)

2025-03-11 15:33:18 -07:00 · 2025-03-11 15:33:18 -07:00 · 553a3d825a
commit 553a3d825a
parent 8e561dced4
10 changed files with 72 additions and 109 deletions
--- a/components/artifact-messages.tsx
+++ b/components/artifact-messages.tsx
@ -1,29 +1,26 @@
 import { PreviewMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Vote } from '@/lib/db/schema';
-import { ChatRequestOptions, Message } from 'ai';
+import { Message } from 'ai';
 import { memo } from 'react';
 import equal from 'fast-deep-equal';
 import { UIArtifact } from './artifact';
+import { UseChatHelpers } from '@ai-sdk/react';

 interface ArtifactMessagesProps {
  chatId: string;
-  isLoading: boolean;
+  status: UseChatHelpers['status'];
  votes: Array<Vote> | undefined;
  messages: Array<Message>;
-  setMessages: (
-    messages: Message[] | ((messages: Message[]) => Message[]),
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
+  setMessages: UseChatHelpers['setMessages'];
+  reload: UseChatHelpers['reload'];
  isReadonly: boolean;
  artifactStatus: UIArtifact['status'];
 }

 function PureArtifactMessages({
  chatId,
-  isLoading,
+  status,
  votes,
  messages,
  setMessages,
@ -43,8 +40,7 @@ function PureArtifactMessages({
          chatId={chatId}
          key={message.id}
          message={message}
-          isLoading={isLoading && index === messages.length - 1}
-          index={index}
+          isLoading={status === 'streaming' && index === messages.length - 1}
          vote={
            votes
              ? votes.find((vote) => vote.messageId === message.id)
@ -74,8 +70,8 @@ function areEqual(
  )
    return true;

-  if (prevProps.isLoading !== nextProps.isLoading) return false;
-  if (prevProps.isLoading && nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
+  if (prevProps.status && nextProps.status) return false;
  if (prevProps.messages.length !== nextProps.messages.length) return false;
  if (!equal(prevProps.votes, nextProps.votes)) return false;

--- a/components/artifact.tsx
+++ b/components/artifact.tsx
@ -1,9 +1,4 @@
-import type {
-  Attachment,
-  ChatRequestOptions,
-  CreateMessage,
-  Message,
-} from 'ai';
+import type { Attachment, Message } from 'ai';
 import { formatDistance } from 'date-fns';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
@ -31,6 +26,7 @@ import { codeArtifact } from '@/artifacts/code/client';
 import { sheetArtifact } from '@/artifacts/sheet/client';
 import { textArtifact } from '@/artifacts/text/client';
 import equal from 'fast-deep-equal';
+import { UseChatHelpers } from '@ai-sdk/react';

 export const artifactDefinitions = [
  textArtifact,
@ -60,7 +56,7 @@ function PureArtifact({
  input,
  setInput,
  handleSubmit,
-  isLoading,
+  status,
  stop,
  attachments,
  setAttachments,
@ -73,27 +69,17 @@ function PureArtifact({
 }: {
  chatId: string;
  input: string;
-  setInput: (input: string) => void;
-  isLoading: boolean;
-  stop: () => void;
+  setInput: UseChatHelpers['setInput'];
+  status: UseChatHelpers['status'];
+  stop: UseChatHelpers['stop'];
  attachments: Array<Attachment>;
  setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
  messages: Array<Message>;
  setMessages: Dispatch<SetStateAction<Array<Message>>>;
  votes: Array<Vote> | undefined;
-  append: (
-    message: Message | CreateMessage,
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
-  handleSubmit: (
-    event?: {
-      preventDefault?: () => void;
-    },
-    chatRequestOptions?: ChatRequestOptions,
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
+  append: UseChatHelpers['append'];
+  handleSubmit: UseChatHelpers['handleSubmit'];
+  reload: UseChatHelpers['reload'];
  isReadonly: boolean;
 }) {
  const { artifact, setArtifact, metadata, setMetadata } = useArtifact();
@ -326,7 +312,7 @@ function PureArtifact({
              <div className="flex flex-col h-full justify-between items-center gap-4">
                <ArtifactMessages
                  chatId={chatId}
-                  isLoading={isLoading}
+                  status={status}
                  votes={votes}
                  messages={messages}
                  setMessages={setMessages}
@ -341,7 +327,7 @@ function PureArtifact({
                    input={input}
                    setInput={setInput}
                    handleSubmit={handleSubmit}
-                    isLoading={isLoading}
+                    status={status}
                    stop={stop}
                    attachments={attachments}
                    setAttachments={setAttachments}
@ -487,7 +473,7 @@ function PureArtifact({
                    isToolbarVisible={isToolbarVisible}
                    setIsToolbarVisible={setIsToolbarVisible}
                    append={append}
-                    isLoading={isLoading}
+                    status={status}
                    stop={stop}
                    setMessages={setMessages}
                    artifactKind={artifact.kind}
@ -513,7 +499,7 @@ function PureArtifact({
 }

 export const Artifact = memo(PureArtifact, (prevProps, nextProps) => {
-  if (prevProps.isLoading !== nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
  if (!equal(prevProps.votes, nextProps.votes)) return false;
  if (prevProps.input !== nextProps.input) return false;
  if (!equal(prevProps.messages, nextProps.messages.length)) return false;
--- a/components/chat.tsx
+++ b/components/chat.tsx
@ -36,7 +36,7 @@ export function Chat({
    input,
    setInput,
    append,
-    isLoading,
+    status,
    stop,
    reload,
  } = useChat({
@ -74,7 +74,7 @@ export function Chat({

        <Messages
          chatId={id}
-          isLoading={isLoading}
+          status={status}
          votes={votes}
          messages={messages}
          setMessages={setMessages}
@ -90,7 +90,7 @@ export function Chat({
              input={input}
              setInput={setInput}
              handleSubmit={handleSubmit}
-              isLoading={isLoading}
+              status={status}
              stop={stop}
              attachments={attachments}
              setAttachments={setAttachments}
@ -107,7 +107,7 @@ export function Chat({
        input={input}
        setInput={setInput}
        handleSubmit={handleSubmit}
-        isLoading={isLoading}
+        status={status}
        stop={stop}
        attachments={attachments}
        setAttachments={setAttachments}
--- a/components/message.tsx
+++ b/components/message.tsx
@ -27,7 +27,6 @@ const PurePreviewMessage = ({
  setMessages,
  reload,
  isReadonly,
-  index,
 }: {
  chatId: string;
  message: Message;
@ -40,7 +39,6 @@ const PurePreviewMessage = ({
    chatRequestOptions?: ChatRequestOptions,
  ) => Promise<string | null | undefined>;
  isReadonly: boolean;
-  index: number;
 }) => {
  const [mode, setMode] = useState<'view' | 'edit'>('view');

--- a/components/messages.tsx
+++ b/components/messages.tsx
@ -1,29 +1,26 @@
-import { ChatRequestOptions, Message } from 'ai';
+import { Message } from 'ai';
 import { PreviewMessage, ThinkingMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Overview } from './overview';
 import { memo } from 'react';
 import { Vote } from '@/lib/db/schema';
 import equal from 'fast-deep-equal';
+import { UseChatHelpers } from '@ai-sdk/react';

 interface MessagesProps {
  chatId: string;
-  isLoading: boolean;
+  status: UseChatHelpers['status'];
  votes: Array<Vote> | undefined;
  messages: Array<Message>;
-  setMessages: (
-    messages: Message[] | ((messages: Message[]) => Message[]),
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
+  setMessages: UseChatHelpers['setMessages'];
+  reload: UseChatHelpers['reload'];
  isReadonly: boolean;
  isArtifactVisible: boolean;
 }

 function PureMessages({
  chatId,
-  isLoading,
+  status,
  votes,
  messages,
  setMessages,
@ -43,10 +40,9 @@ function PureMessages({
      {messages.map((message, index) => (
        <PreviewMessage
          key={message.id}
-          index={index}
          chatId={chatId}
          message={message}
-          isLoading={isLoading && messages.length - 1 === index}
+          isLoading={status === 'streaming' && messages.length - 1 === index}
          vote={
            votes
              ? votes.find((vote) => vote.messageId === message.id)
@ -58,7 +54,7 @@ function PureMessages({
        />
      ))}

-      {isLoading &&
+      {status === 'submitted' &&
        messages.length > 0 &&
        messages[messages.length - 1].role === 'user' && <ThinkingMessage />}

@ -73,8 +69,8 @@ function PureMessages({
 export const Messages = memo(PureMessages, (prevProps, nextProps) => {
  if (prevProps.isArtifactVisible && nextProps.isArtifactVisible) return true;

-  if (prevProps.isLoading !== nextProps.isLoading) return false;
-  if (prevProps.isLoading && nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
+  if (prevProps.status && nextProps.status) return false;
  if (prevProps.messages.length !== nextProps.messages.length) return false;
  if (!equal(prevProps.messages, nextProps.messages)) return false;
  if (!equal(prevProps.votes, nextProps.votes)) return false;
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@ -29,12 +29,13 @@ import { Button } from './ui/button';
 import { Textarea } from './ui/textarea';
 import { SuggestedActions } from './suggested-actions';
 import equal from 'fast-deep-equal';
+import { UseChatHelpers, UseChatOptions } from '@ai-sdk/react';

 function PureMultimodalInput({
  chatId,
  input,
  setInput,
-  isLoading,
+  status,
  stop,
  attachments,
  setAttachments,
@ -45,24 +46,16 @@ function PureMultimodalInput({
  className,
 }: {
  chatId: string;
-  input: string;
-  setInput: (value: string) => void;
-  isLoading: boolean;
+  input: UseChatHelpers['input'];
+  setInput: UseChatHelpers['setInput'];
+  status: UseChatHelpers['status'];
  stop: () => void;
  attachments: Array<Attachment>;
  setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
  messages: Array<Message>;
  setMessages: Dispatch<SetStateAction<Array<Message>>>;
-  append: (
-    message: Message | CreateMessage,
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
-  handleSubmit: (
-    event?: {
-      preventDefault?: () => void;
-    },
-    chatRequestOptions?: ChatRequestOptions,
-  ) => void;
+  append: UseChatHelpers['append'];
+  handleSubmit: UseChatHelpers['handleSubmit'];
  className?: string;
 }) {
  const textareaRef = useRef<HTMLTextAreaElement>(null);
@ -253,7 +246,7 @@ function PureMultimodalInput({
          ) {
            event.preventDefault();

-            if (isLoading) {
+            if (status !== 'ready') {
              toast.error('Please wait for the model to finish its response!');
            } else {
              submitForm();
@ -263,11 +256,11 @@ function PureMultimodalInput({
      />

      <div className="absolute bottom-0 p-2 w-fit flex flex-row justify-start">
-        <AttachmentsButton fileInputRef={fileInputRef} isLoading={isLoading} />
+        <AttachmentsButton fileInputRef={fileInputRef} status={status} />
      </div>

      <div className="absolute bottom-0 right-0 p-2 w-fit flex flex-row justify-end">
-        {isLoading ? (
+        {status === 'submitted' ? (
          <StopButton stop={stop} setMessages={setMessages} />
        ) : (
          <SendButton
@ -285,7 +278,7 @@ export const MultimodalInput = memo(
  PureMultimodalInput,
  (prevProps, nextProps) => {
    if (prevProps.input !== nextProps.input) return false;
-    if (prevProps.isLoading !== nextProps.isLoading) return false;
+    if (prevProps.status !== nextProps.status) return false;
    if (!equal(prevProps.attachments, nextProps.attachments)) return false;

    return true;
@ -294,10 +287,10 @@ export const MultimodalInput = memo(

 function PureAttachmentsButton({
  fileInputRef,
-  isLoading,
+  status,
 }: {
  fileInputRef: React.MutableRefObject<HTMLInputElement | null>;
-  isLoading: boolean;
+  status: UseChatHelpers['status'];
 }) {
  return (
    <Button
@ -307,7 +300,7 @@ function PureAttachmentsButton({
        event.preventDefault();
        fileInputRef.current?.click();
      }}
-      disabled={isLoading}
+      disabled={status !== 'ready'}
      variant="ghost"
    >
      <PaperclipIcon size={14} />
--- a/components/toolbar.tsx
+++ b/components/toolbar.tsx
@ -319,19 +319,16 @@ const PureToolbar = ({
  isToolbarVisible,
  setIsToolbarVisible,
  append,
-  isLoading,
+  status,
  stop,
  setMessages,
  artifactKind,
 }: {
  isToolbarVisible: boolean;
  setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
-  isLoading: boolean;
-  append: (
-    message: Message | CreateMessage,
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
-  stop: () => void;
+  status: UseChatHelpers['status'];
+  append: UseChatHelpers['append'];
+  stop: UseChatHelpers['stop'];
  setMessages: Dispatch<SetStateAction<Message[]>>;
  artifactKind: ArtifactKind;
 }) => {
@ -372,10 +369,10 @@ const PureToolbar = ({
  }, []);

  useEffect(() => {
-    if (isLoading) {
+    if (status === 'streaming') {
      setIsToolbarVisible(false);
    }
-  }, [isLoading, setIsToolbarVisible]);
+  }, [status, setIsToolbarVisible]);

  const artifactDefinition = artifactDefinitions.find(
    (definition) => definition.kind === artifactKind,
@ -418,13 +415,13 @@ const PureToolbar = ({
        exit={{ opacity: 0, y: -20, transition: { duration: 0.1 } }}
        transition={{ type: 'spring', stiffness: 300, damping: 25 }}
        onHoverStart={() => {
-          if (isLoading) return;
+          if (status === 'streaming') return;

          cancelCloseTimer();
          setIsToolbarVisible(true);
        }}
        onHoverEnd={() => {
-          if (isLoading) return;
+          if (status === 'streaming') return;

          startCloseTimer();
        }}
@ -436,7 +433,7 @@ const PureToolbar = ({
        }}
        ref={toolbarRef}
      >
-        {isLoading ? (
+        {status === 'streaming' ? (
          <motion.div
            key="stop-icon"
            initial={{ scale: 1 }}
@ -475,7 +472,7 @@ const PureToolbar = ({
 };

 export const Toolbar = memo(PureToolbar, (prevProps, nextProps) => {
-  if (prevProps.isLoading !== nextProps.isLoading) return false;
+  if (prevProps.status !== nextProps.status) return false;
  if (prevProps.isToolbarVisible !== nextProps.isToolbarVisible) return false;
  if (prevProps.artifactKind !== nextProps.artifactKind) return false;