Merge pull request #4150 from janhq/fix/performance-issue-on-chat-generation

fix: convo component re-render every single token yield

web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx

@@ -1,6 +1,6 @@
-import { memo } from 'react'
+import { memo, useEffect, useState } from 'react'
 
-import { MessageStatus } from '@janhq/core'
+import { MessageStatus, ThreadMessage } from '@janhq/core'
 
 import { useAtomValue } from 'jotai'
 
@@ -17,33 +17,63 @@ import EmptyThread from './EmptyThread'
 
 import { getCurrentChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
 
-const ChatBody = () => {
+const ChatConfigurator = memo(() => {
   const messages = useAtomValue(getCurrentChatMessagesAtom)
 
+  const [current, setCurrent] = useState<ThreadMessage[]>([])
+
+  const isMessagesIdentificial = (
+    arr1: ThreadMessage[],
+    arr2: ThreadMessage[]
+  ): boolean => {
+    if (arr1.length !== arr2.length) return false
+    return arr1.every((item, index) => item.id === arr2[index].id)
+  }
+
+  useEffect(() => {
+    if (
+      messages.length !== current.length ||
+      !isMessagesIdentificial(messages, current)
+    ) {
+      setCurrent(messages)
+    }
+  }, [messages, current])
+
   const loadModelError = useAtomValue(loadModelErrorAtom)
 
   if (!messages.length) return <EmptyThread />
+  return (
+    <div className="flex h-full w-full flex-col">
+      <ChatBody loadModelError={loadModelError} messages={current} />
+    </div>
+  )
+})
 
+const ChatBody = memo(
+  ({
+    messages,
+    loadModelError,
+  }: {
+    messages: ThreadMessage[]
+    loadModelError?: string
+  }) => {
     return (
       <ListContainer>
         {messages.map((message, index) => (
           <div key={message.id}>
-          {message.status !== MessageStatus.Error &&
-            message.content.length > 0 && (
-              <ChatItem {...message} key={message.id} />
-            )}
-
-          {!loadModelError &&
-            index === messages.length - 1 &&
-            message.status !== MessageStatus.Pending &&
-            message.status !== MessageStatus.Ready && (
-              <ErrorMessage message={message} />
-            )}
+            <ChatItem
+              {...message}
+              key={message.id}
+              loadModelError={loadModelError}
+              isCurrentMessage={index === messages.length - 1}
+            />
           </div>
         ))}
+
         {loadModelError && <LoadModelError />}
       </ListContainer>
     )
-}
+  }
+)
 
-export default memo(ChatBody)
+export default memo(ChatConfigurator)

web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx

@@ -1,15 +1,67 @@
-import React, { forwardRef } from 'react'
+import React, { forwardRef, useEffect, useState } from 'react'
 
-import { ThreadMessage } from '@janhq/core'
+import {
+  events,
+  MessageEvent,
+  MessageStatus,
+  ThreadContent,
+  ThreadMessage,
+} from '@janhq/core'
+
+import ErrorMessage from '@/containers/ErrorMessage'
 
 import SimpleTextMessage from '../SimpleTextMessage'
 
 type Ref = HTMLDivElement
 
-const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => (
+type Props = {
+  loadModelError?: string
+  isCurrentMessage?: boolean
+} & ThreadMessage
+
+const ChatItem = forwardRef<Ref, Props>((message, ref) => {
+  const [content, setContent] = useState<ThreadContent[]>(message.content)
+  const [status, setStatus] = useState<MessageStatus>(message.status)
+  const [errorMessage, setErrorMessage] = useState<ThreadMessage | undefined>(
+    message.isCurrentMessage && message.status === MessageStatus.Error
+      ? message
+      : undefined
+  )
+
+  function onMessageUpdate(data: ThreadMessage) {
+    if (data.id === message.id) {
+      setContent(data.content)
+      if (data.status !== status) setStatus(data.status)
+      if (data.status === MessageStatus.Error && message.isCurrentMessage)
+        setErrorMessage(data)
+    }
+  }
+
+  useEffect(() => {
+    if (!message.isCurrentMessage && errorMessage) setErrorMessage(undefined)
+  }, [message, errorMessage])
+
+  useEffect(() => {
+    if (message.status === MessageStatus.Pending)
+      events.on(MessageEvent.OnMessageUpdate, onMessageUpdate)
+    return () => {
+      events.off(MessageEvent.OnMessageUpdate, onMessageUpdate)
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  return (
+    <>
+      {status !== MessageStatus.Error && content.length > 0 && (
         <div ref={ref} className="relative">
-    <SimpleTextMessage {...message} />
+          <SimpleTextMessage {...message} content={content} status={status} />
         </div>
-))
+      )}
+      {errorMessage && !message.loadModelError && (
+        <ErrorMessage message={errorMessage} />
+      )}
+    </>
+  )
+})
 
 export default ChatItem