From 08cba390a71756ba7ac44d4816a1629d9946ecaf Mon Sep 17 00:00:00 2001 From: chaos <1585767724@qq.com> Date: Wed, 4 Sep 2024 14:27:51 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BD=BF=E7=94=A8useMemo=E4=BC=98=E5=8C=96send?= =?UTF-8?q?Message=E4=B8=AD=E7=9A=84=E6=B6=88=E6=81=AF=E7=94=9F=E6=88=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/hooks/useChatList.ts | 194 ++++++++++++++++++++++----------------- 1 file changed, 109 insertions(+), 85 deletions(-) diff --git a/src/hooks/useChatList.ts b/src/hooks/useChatList.ts index ee41918..8d250fa 100644 --- a/src/hooks/useChatList.ts +++ b/src/hooks/useChatList.ts @@ -5,7 +5,7 @@ import { ModelConfig } from '@/types/config'; import { processSSE } from '@/utils/fetch'; import { genUUID } from '@/utils/uuid'; import { useMergedState } from 'rc-util'; -import React, { useEffect, useMemo, useRef } from 'react'; +import React, { useCallback, useEffect, useMemo, useRef } from 'react'; import { useRefFunction } from './useRefFunction'; export const initialModelConfig: ModelConfig = { @@ -129,6 +129,7 @@ export const useChatList = (props: ProChatUIUseListChatProps) => { value: props.loading, }); const [isLoadingMessage, setIsLoadingMessage] = useMergedState(false); + /** * Fetches the chat list using the provided request function. * If the request function is not provided, it sets the loading state to false and returns. @@ -180,102 +181,125 @@ export const useChatList = (props: ProChatUIUseListChatProps) => { * @param {string} message - The message to send. * @returns {Promise} - A promise that resolves when the message is sent. */ + const getLastLoadingMessage = useCallback(() => { + return chatListRef.current[chatListRef.current.length - 1]; + }, [chatListRef]); + const sendMessage = useRefFunction(async (message: string | Partial) => { controller.current = new AbortController(); - setChatList((prevState) => [ - ...prevState, - genMessageRecord( - typeof message === 'string' - ? { content: message } - : { - ...message, - content: message.content, - }, - 'user', - ), - ]); + const userMessage = genMessageRecord( + typeof message === 'string' ? { content: message } : { ...message, content: message.content }, + 'user', + ); + chatList.push(userMessage); + setChatList([...chatList]); if (!props?.sendMessageRequest) return; setIsLoadingMessage(true); - setChatList((prevState) => [ - ...prevState, - genMessageRecord( - { - content: LOADING_FLAT, - }, - 'assistant', - ), - ]); - - const res = (await Promise.race([ - props.sendMessageRequest?.(chatListRef.current), - new Promise((_, reject) => { - controller.current?.signal.addEventListener('abort', () => { - reject(); - }); - }), - ])) as Response | ChatMessage; - - if (res instanceof Response) { - processSSE(res, { - signal: controller.current.signal, - onFinish: async () => { - setChatList((prevState) => { - const updatedList = [...prevState]; - updatedList[updatedList.length - 1].isFinished = true; - return updatedList; + const loadingMessage = genMessageRecord({ content: LOADING_FLAT }, 'assistant'); + chatList.push(loadingMessage); + setChatList([...chatList]); + + try { + const res = (await Promise.race([ + props.sendMessageRequest?.([chatList[chatList.length - 2]]), + new Promise((_, reject) => { + controller.current.signal.addEventListener('abort', () => { + reject(); }); - setIsLoadingMessage(false); - }, - onMessageHandle: async (text, res, type) => { - if (type === 'done' || controller.current?.signal.aborted) { + }), + ])) as Response | ChatMessage; + + if (res instanceof Response) { + await processSSE(res, { + signal: controller.current.signal, + onFinish: async () => { + const message = getLastLoadingMessage(); + if (!message) return; + setChatList((prev) => { + message.isFinished = true; + return [...prev]; + }); setIsLoadingMessage(false); - setChatList((prevState) => { - const updatedList = [...prevState]; - updatedList[updatedList.length - 1].isFinished = true; - return updatedList; + }, + onMessageHandle: async (text, res, type) => { + const loadingMessage = getLastLoadingMessage(); + if (!loadingMessage) return; + + if (type === 'done' || controller.current.signal.aborted) { + loadingMessage.isFinished = true; + setChatList((prev) => [...prev]); + setIsLoadingMessage(false); + return; + } + + const content = + loadingMessage.content === LOADING_FLAT ? text : loadingMessage.content + text; + const updatedMessage = { + ...loadingMessage, + updateAt: Date.now(), + originContent: text, + isFinished: false, + content, + }; + + const transformMessage = await props.transformToChatMessage?.(updatedMessage, { + preContent: loadingMessage.content === LOADING_FLAT ? '' : loadingMessage.content, + currentContent: text, }); - return; - } - - setChatList((prevState) => { - const updatedList = [...prevState]; - const currentMessage = updatedList[updatedList.length - 1]; - currentMessage.content = - currentMessage.content === LOADING_FLAT ? text : currentMessage.content + text; - currentMessage.updateAt = Date.now(); - return updatedList; - }); - }, - onErrorHandle: async (error) => { - const content = error.message; - setChatList((prevState) => { - const updatedList = [...prevState]; - const currentMessage = updatedList[updatedList.length - 1]; - currentMessage.content = content; - currentMessage.originContent = content; - currentMessage.updateAt = Date.now(); - currentMessage.isFinished = true; - return updatedList; + + setChatList((prev) => { + const newList = [...prev]; + newList[newList.length - 1] = transformMessage || updatedMessage; + return newList; + }); + }, + onErrorHandle: async (error) => { + const loadingMessage = getLastLoadingMessage(); + if (loadingMessage) { + const content = error.message; + const errorMessage = { + ...loadingMessage, + content, + originContent: content, + updateAt: Date.now(), + isFinished: true, + }; + + const transformMessage = await props.transformToChatMessage?.(errorMessage, { + preContent: loadingMessage.content, + currentContent: content, + }); + + setChatList((prev) => [...prev, transformMessage || errorMessage]); + } + setIsLoadingMessage(false); + }, + }); + } else { + const loadingMessage = getLastLoadingMessage(); + if (loadingMessage) { + const updatedMessage = { + ...loadingMessage, + ...res, + updateAt: Date.now(), + }; + + const transformMessage = await props.transformToChatMessage?.(updatedMessage, { + preContent: loadingMessage.content, + currentContent: updatedMessage.originContent, }); - setIsLoadingMessage(false); - }, - }); - } else { - setChatList((prevState) => { - const updatedList = [...prevState]; - const currentMessage = { - ...updatedList[updatedList.length - 1], - ...res, - updateAt: Date.now(), - }; - updatedList[updatedList.length - 1] = currentMessage; - return updatedList; - }); - setIsLoadingMessage(false); + + setChatList((prev) => [...prev, transformMessage || updatedMessage]); + } + setIsLoadingMessage(false); + } + } catch (error) { + console.error('Error sending message:', error); } }); + /** * Stops the generation of messages. */