Skip to content

Commit

Permalink
fix: convo component re-render every single token yield
Browse files Browse the repository at this point in the history
  • Loading branch information
louis-jan committed Nov 28, 2024
1 parent cd1a274 commit 2e35b2f
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 29 deletions.
78 changes: 56 additions & 22 deletions web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx
Original file line number Diff line number Diff line change
@@ -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'

Expand All @@ -17,33 +17,67 @@ 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 (
<ListContainer>
{messages.map((message, index) => (
<div key={message.id}>
{message.status !== MessageStatus.Error &&
message.content.length > 0 && (
<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 && (
<ChatItem {...message} key={message.id} />
)}

{!loadModelError &&
index === messages.length - 1 &&
message.status !== MessageStatus.Pending &&
message.status !== MessageStatus.Ready && (
<ErrorMessage message={message} />
)}
</div>
))}
{loadModelError && <LoadModelError />}
</ListContainer>
)
}
{!loadModelError &&
index === messages.length - 1 &&
message.status !== MessageStatus.Pending &&
message.status !== MessageStatus.Ready && (
<ErrorMessage message={message} />
)}
</div>
))}

{loadModelError && <LoadModelError />}
</ListContainer>
)
}
)

export default memo(ChatBody)
export default memo(ChatConfigurator)
41 changes: 34 additions & 7 deletions web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
Original file line number Diff line number Diff line change
@@ -1,15 +1,42 @@
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 SimpleTextMessage from '../SimpleTextMessage'

type Ref = HTMLDivElement

const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => (
<div ref={ref} className="relative">
<SimpleTextMessage {...message} />
</div>
))
const ChatItem = forwardRef<Ref, ThreadMessage>((message, ref) => {
const [content, setContent] = useState<ThreadContent[]>(message.content)
const [status, setStatus] = useState<MessageStatus>(message.status)

function onMessageUpdate(data: ThreadMessage) {
if (data.id === message.id) {
setContent(data.content)
if (data.status !== status) setStatus(data.status)
}
}

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 (
<div ref={ref} className="relative">
<SimpleTextMessage {...message} content={content} status={status} />
</div>
)
})

export default ChatItem

0 comments on commit 2e35b2f

Please sign in to comment.