import { useCallback, useEffect, useRef, useState, } from 'react' import { useTranslation } from 'react-i18next' import { produce, setAutoFreeze } from 'immer' import { useWorkflowRun } from '../../hooks' import type { ChatItem } from '@/app/components/base/chat/types' import { useToastContext } from '@/app/components/base/toast' import { TransferMethod } from '@/types/app' import type { VisionFile } from '@/types/app' export const useChat = ( prevChatList?: ChatItem[], ) => { const { t } = useTranslation() const { notify } = useToastContext() const { handleRun } = useWorkflowRun() const hasStopResponded = useRef(false) const connversationId = useRef('') const taskIdRef = useRef('') const [chatList, setChatList] = useState(prevChatList || []) const chatListRef = useRef(prevChatList || []) const [isResponding, setIsResponding] = useState(false) const isRespondingRef = useRef(false) const [suggestedQuestions, setSuggestQuestions] = useState([]) const stopAbortControllerRef = useRef(null) const suggestedQuestionsAbortControllerRef = useRef(null) useEffect(() => { setAutoFreeze(false) return () => { setAutoFreeze(true) } }, []) const handleUpdateChatList = useCallback((newChatList: ChatItem[]) => { setChatList(newChatList) chatListRef.current = newChatList }, []) const handleResponding = useCallback((isResponding: boolean) => { setIsResponding(isResponding) isRespondingRef.current = isResponding }, []) const handleStop = useCallback(() => { hasStopResponded.current = true handleResponding(false) }, [handleResponding]) const updateCurrentQA = useCallback(({ responseItem, questionId, placeholderAnswerId, questionItem, }: { responseItem: ChatItem questionId: string placeholderAnswerId: string questionItem: ChatItem }) => { const newListWithAnswer = produce( chatListRef.current.filter(item => item.id !== responseItem.id && item.id !== placeholderAnswerId), (draft) => { if (!draft.find(item => item.id === questionId)) draft.push({ ...questionItem }) draft.push({ ...responseItem }) }) handleUpdateChatList(newListWithAnswer) }, [handleUpdateChatList]) const handleSend = useCallback((params: any) => { if (isRespondingRef.current) { notify({ type: 'info', message: t('appDebug.errorMessage.waitForResponse') }) return false } const questionId = `question-${Date.now()}` const questionItem = { id: questionId, content: params.query, isAnswer: false, message_files: params.files, } const placeholderAnswerId = `answer-placeholder-${Date.now()}` const placeholderAnswerItem = { id: placeholderAnswerId, content: '', isAnswer: true, } const newList = [...chatListRef.current, questionItem, placeholderAnswerItem] handleUpdateChatList(newList) // answer const responseItem: ChatItem = { id: `${Date.now()}`, content: '', agent_thoughts: [], message_files: [], isAnswer: true, } handleResponding(true) const bodyParams = { conversation_id: connversationId.current, ...params, } if (bodyParams?.files?.length) { bodyParams.files = bodyParams.files.map((item: VisionFile) => { if (item.transfer_method === TransferMethod.local_file) { return { ...item, url: '', } } return item }) } let hasSetResponseId = false handleRun( params, { onData: (message: string, isFirstMessage: boolean, { conversationId: newConversationId, messageId, taskId }: any) => { responseItem.content = responseItem.content + message if (messageId && !hasSetResponseId) { responseItem.id = messageId hasSetResponseId = true } if (isFirstMessage && newConversationId) connversationId.current = newConversationId taskIdRef.current = taskId if (messageId) responseItem.id = messageId updateCurrentQA({ responseItem, questionId, placeholderAnswerId, questionItem, }) }, async onCompleted(hasError?: boolean) { handleResponding(false) }, onMessageEnd: (messageEnd) => { responseItem.citation = messageEnd.metadata?.retriever_resources || [] const newListWithAnswer = produce( chatListRef.current.filter(item => item.id !== responseItem.id && item.id !== placeholderAnswerId), (draft) => { if (!draft.find(item => item.id === questionId)) draft.push({ ...questionItem }) draft.push({ ...responseItem }) }) handleUpdateChatList(newListWithAnswer) }, onMessageReplace: (messageReplace) => { responseItem.content = messageReplace.answer }, onError() { handleResponding(false) const newChatList = produce(chatListRef.current, (draft) => { draft.splice(draft.findIndex(item => item.id === placeholderAnswerId), 1) }) handleUpdateChatList(newChatList) }, }, ) }, [handleRun, handleResponding, handleUpdateChatList, notify, t, updateCurrentQA]) return { conversationId: connversationId.current, chatList, handleSend, handleStop, isResponding, suggestedQuestions, } }