feat: switch to chat messages before regenerated (#11301)

Co-authored-by: zuodongxu <192560071+zuodongxu@users.noreply.github.com>

api/controllers/console/explore/message.py

@@ -50,7 +50,7 @@ class MessageListApi(InstalledAppResource):
 
         try:
             return MessageService.pagination_by_first_id(
-                app_model, current_user, args["conversation_id"], args["first_id"], args["limit"], "desc"
+                app_model, current_user, args["conversation_id"], args["first_id"], args["limit"]
             )
         except services.errors.conversation.ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")

api/controllers/web/message.py

@@ -91,7 +91,7 @@ class MessageListApi(WebApiResource):
 
         try:
             return MessageService.pagination_by_first_id(
-                app_model, end_user, args["conversation_id"], args["first_id"], args["limit"], "desc"
+                app_model, end_user, args["conversation_id"], args["first_id"], args["limit"]
             )
         except services.errors.conversation.ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")

web/app/components/app/configuration/debug/debug-with-multiple-model/chat-item.tsx

@@ -67,7 +67,6 @@ const ChatItem: FC<ChatItemProps> = ({
   }, [modelConfig.configs.prompt_variables])
   const {
     chatList,
-    chatListRef,
     isResponding,
     handleSend,
     suggestedQuestions,
@@ -102,7 +101,7 @@ const ChatItem: FC<ChatItemProps> = ({
       query: message,
       inputs,
       model_config: configData,
-      parent_message_id: getLastAnswer(chatListRef.current)?.id || null,
+      parent_message_id: getLastAnswer(chatList)?.id || null,
     }
 
     if ((config.file_upload as any).enabled && files?.length && supportVision)
@@ -116,7 +115,7 @@ const ChatItem: FC<ChatItemProps> = ({
         onGetSuggestedQuestions: (responseItemId, getAbortController) => fetchSuggestedQuestions(appId, responseItemId, getAbortController),
       },
     )
-  }, [appId, config, handleSend, inputs, modelAndParameter, textGenerationModelList, chatListRef])
+  }, [appId, chatList, config, handleSend, inputs, modelAndParameter.model, modelAndParameter.parameters, modelAndParameter.provider, textGenerationModelList])
 
   const { eventEmitter } = useEventEmitterContextContext()
   eventEmitter?.useSubscription((v: any) => {

web/app/components/app/configuration/debug/debug-with-single-model/index.tsx

@@ -12,7 +12,7 @@ import {
 import Chat from '@/app/components/base/chat/chat'
 import { useChat } from '@/app/components/base/chat/chat/hooks'
 import { useDebugConfigurationContext } from '@/context/debug-configuration'
-import type { ChatConfig, ChatItem, OnSend } from '@/app/components/base/chat/types'
+import type { ChatConfig, ChatItem, ChatItemInTree, OnSend } from '@/app/components/base/chat/types'
 import { useProviderContext } from '@/context/provider-context'
 import {
   fetchConversationMessages,
@@ -24,7 +24,7 @@ import { useAppContext } from '@/context/app-context'
 import { ModelFeatureEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
 import { useStore as useAppStore } from '@/app/components/app/store'
 import { useFeatures } from '@/app/components/base/features/hooks'
-import { getLastAnswer } from '@/app/components/base/chat/utils'
+import { getLastAnswer, isValidGeneratedAnswer } from '@/app/components/base/chat/utils'
 import type { InputForm } from '@/app/components/base/chat/chat/type'
 
 type DebugWithSingleModelProps = {
@@ -68,12 +68,11 @@ const DebugWithSingleModel = forwardRef<DebugWithSingleModelRefType, DebugWithSi
   }, [modelConfig.configs.prompt_variables])
   const {
     chatList,
-    chatListRef,
+    setTargetMessageId,
     isResponding,
     handleSend,
     suggestedQuestions,
     handleStop,
-    handleUpdateChatList,
     handleRestart,
     handleAnnotationAdded,
     handleAnnotationEdited,
@@ -89,7 +88,7 @@ const DebugWithSingleModel = forwardRef<DebugWithSingleModelRefType, DebugWithSi
   )
   useFormattingChangedSubscription(chatList)
 
-  const doSend: OnSend = useCallback((message, files, last_answer) => {
+  const doSend: OnSend = useCallback((message, files, isRegenerate = false, parentAnswer: ChatItem | null = null) => {
     if (checkCanSend && !checkCanSend())
       return
     const currentProvider = textGenerationModelList.find(item => item.provider === modelConfig.provider)
@@ -110,7 +109,7 @@ const DebugWithSingleModel = forwardRef<DebugWithSingleModelRefType, DebugWithSi
       query: message,
       inputs,
       model_config: configData,
-      parent_message_id: last_answer?.id || getLastAnswer(chatListRef.current)?.id || null,
+      parent_message_id: (isRegenerate ? parentAnswer?.id : getLastAnswer(chatList)?.id) || null,
     }
 
     if ((config.file_upload as any)?.enabled && files?.length && supportVision)
@@ -124,23 +123,13 @@ const DebugWithSingleModel = forwardRef<DebugWithSingleModelRefType, DebugWithSi
         onGetSuggestedQuestions: (responseItemId, getAbortController) => fetchSuggestedQuestions(appId, responseItemId, getAbortController),
       },
     )
-  }, [chatListRef, appId, checkCanSend, completionParams, config, handleSend, inputs, modelConfig, textGenerationModelList])
+  }, [appId, chatList, checkCanSend, completionParams, config, handleSend, inputs, modelConfig.mode, modelConfig.model_id, modelConfig.provider, textGenerationModelList])
 
-  const doRegenerate = useCallback((chatItem: ChatItem) => {
-    const index = chatList.findIndex(item => item.id === chatItem.id)
-    if (index === -1)
-      return
-
-    const prevMessages = chatList.slice(0, index)
-    const question = prevMessages.pop()
-    const lastAnswer = getLastAnswer(prevMessages)
-
-    if (!question)
-      return
-
-    handleUpdateChatList(prevMessages)
-    doSend(question.content, question.message_files, lastAnswer)
-  }, [chatList, handleUpdateChatList, doSend])
+  const doRegenerate = useCallback((chatItem: ChatItemInTree) => {
+    const question = chatList.find(item => item.id === chatItem.parentMessageId)!
+    const parentAnswer = chatList.find(item => item.id === question.parentMessageId)
+    doSend(question.content, question.message_files, true, isValidGeneratedAnswer(parentAnswer) ? parentAnswer : null)
+  }, [chatList, doSend])
 
   const allToolIcons = useMemo(() => {
     const icons: Record<string, any> = {}
@@ -173,6 +162,7 @@ const DebugWithSingleModel = forwardRef<DebugWithSingleModelRefType, DebugWithSi
       inputs={inputs}
       inputsForm={inputsForm}
       onRegenerate={doRegenerate}
+      switchSibling={siblingMessageId => setTargetMessageId(siblingMessageId)}
       onStopResponding={handleStop}
       showPromptLog
       questionIcon={<Avatar avatar={userProfile.avatar_url} name={userProfile.name} size={40} />}

web/app/components/base/chat/chat-with-history/chat-wrapper.tsx

@@ -3,10 +3,11 @@ import Chat from '../chat'
 import type {
   ChatConfig,
   ChatItem,
+  ChatItemInTree,
   OnSend,
 } from '../types'
 import { useChat } from '../chat/hooks'
-import { getLastAnswer } from '../utils'
+import { getLastAnswer, isValidGeneratedAnswer } from '../utils'
 import { useChatWithHistoryContext } from './context'
 import Header from './header'
 import ConfigPanel from './config-panel'
@@ -20,7 +21,7 @@ import AnswerIcon from '@/app/components/base/answer-icon'
 const ChatWrapper = () => {
   const {
     appParams,
-    appPrevChatList,
+    appPrevChatTree,
     currentConversationId,
     currentConversationItem,
     inputsForms,
@@ -50,8 +51,7 @@ const ChatWrapper = () => {
   }, [appParams, currentConversationItem?.introduction, currentConversationId])
   const {
     chatList,
-    chatListRef,
-    handleUpdateChatList,
+    setTargetMessageId,
     handleSend,
     handleStop,
     isResponding,
@@ -62,7 +62,7 @@ const ChatWrapper = () => {
       inputs: (currentConversationId ? currentConversationItem?.inputs : newConversationInputs) as any,
       inputsForm: inputsForms,
     },
-    appPrevChatList,
+    appPrevChatTree,
     taskId => stopChatMessageResponding('', taskId, isInstalledApp, appId),
   )
 
@@ -72,13 +72,13 @@ const ChatWrapper = () => {
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
-  const doSend: OnSend = useCallback((message, files, last_answer) => {
+  const doSend: OnSend = useCallback((message, files, isRegenerate = false, parentAnswer: ChatItem | null = null) => {
     const data: any = {
       query: message,
       files,
       inputs: currentConversationId ? currentConversationItem?.inputs : newConversationInputs,
       conversation_id: currentConversationId,
-      parent_message_id: last_answer?.id || getLastAnswer(chatListRef.current)?.id || null,
+      parent_message_id: (isRegenerate ? parentAnswer?.id : getLastAnswer(chatList)?.id) || null,
     }
 
     handleSend(
@@ -91,31 +91,21 @@ const ChatWrapper = () => {
       },
     )
   }, [
-    chatListRef,
+    chatList,
+    handleNewConversationCompleted,
+    handleSend,
     currentConversationId,
     currentConversationItem,
-    handleSend,
     newConversationInputs,
-    handleNewConversationCompleted,
     isInstalledApp,
     appId,
   ])
 
-  const doRegenerate = useCallback((chatItem: ChatItem) => {
-    const index = chatList.findIndex(item => item.id === chatItem.id)
-    if (index === -1)
-      return
-
-    const prevMessages = chatList.slice(0, index)
-    const question = prevMessages.pop()
-    const lastAnswer = getLastAnswer(prevMessages)
-
-    if (!question)
-      return
-
-    handleUpdateChatList(prevMessages)
-    doSend(question.content, question.message_files, lastAnswer)
-  }, [chatList, handleUpdateChatList, doSend])
+  const doRegenerate = useCallback((chatItem: ChatItemInTree) => {
+    const question = chatList.find(item => item.id === chatItem.parentMessageId)!
+    const parentAnswer = chatList.find(item => item.id === question.parentMessageId)
+    doSend(question.content, question.message_files, true, isValidGeneratedAnswer(parentAnswer) ? parentAnswer : null)
+  }, [chatList, doSend])
 
   const chatNode = useMemo(() => {
     if (inputsForms.length) {
@@ -187,6 +177,7 @@ const ChatWrapper = () => {
         answerIcon={answerIcon}
         hideProcessDetail
         themeBuilder={themeBuilder}
+        switchSibling={siblingMessageId => setTargetMessageId(siblingMessageId)}
       />
     </div>
   )

web/app/components/base/chat/chat-with-history/context.tsx

@@ -5,7 +5,7 @@ import { createContext, useContext } from 'use-context-selector'
 import type {
   Callback,
   ChatConfig,
-  ChatItem,
+  ChatItemInTree,
   Feedback,
 } from '../types'
 import type { ThemeBuilder } from '../embedded-chatbot/theme/theme-context'
@@ -25,7 +25,7 @@ export type ChatWithHistoryContextValue = {
   appChatListDataLoading?: boolean
   currentConversationId: string
   currentConversationItem?: ConversationItem
-  appPrevChatList: ChatItem[]
+  appPrevChatTree: ChatItemInTree[]
   pinnedConversationList: AppConversationData['data']
   conversationList: AppConversationData['data']
   showConfigPanelBeforeChat: boolean
@@ -53,7 +53,7 @@ export type ChatWithHistoryContextValue = {
 
 export const ChatWithHistoryContext = createContext<ChatWithHistoryContextValue>({
   currentConversationId: '',
-  appPrevChatList: [],
+  appPrevChatTree: [],
   pinnedConversationList: [],
   conversationList: [],
   showConfigPanelBeforeChat: false,

web/app/components/base/chat/chat-with-history/hooks.tsx

@@ -12,10 +12,13 @@ import produce from 'immer'
 import type {
   Callback,
   ChatConfig,
+  ChatItem,
   Feedback,
 } from '../types'
 import { CONVERSATION_ID_INFO } from '../constants'
-import { getPrevChatList } from '../utils'
+import { buildChatItemTree } from '../utils'
+import { addFileInfos, sortAgentSorts } from '../../../tools/utils'
+import { getProcessedFilesFromResponse } from '@/app/components/base/file-uploader/utils'
 import {
   delConversation,
   fetchAppInfo,
@@ -40,6 +43,32 @@ import { useAppFavicon } from '@/hooks/use-app-favicon'
 import { InputVarType } from '@/app/components/workflow/types'
 import { TransferMethod } from '@/types/app'
 
+function getFormattedChatList(messages: any[]) {
+  const newChatList: ChatItem[] = []
+  messages.forEach((item) => {
+    const questionFiles = item.message_files?.filter((file: any) => file.belongs_to === 'user') || []
+    newChatList.push({
+      id: `question-${item.id}`,
+      content: item.query,
+      isAnswer: false,
+      message_files: getProcessedFilesFromResponse(questionFiles.map((item: any) => ({ ...item, related_id: item.id }))),
+      parentMessageId: item.parent_message_id || undefined,
+    })
+    const answerFiles = item.message_files?.filter((file: any) => file.belongs_to === 'assistant') || []
+    newChatList.push({
+      id: item.id,
+      content: item.answer,
+      agent_thoughts: addFileInfos(item.agent_thoughts ? sortAgentSorts(item.agent_thoughts) : item.agent_thoughts, item.message_files),
+      feedback: item.feedback,
+      isAnswer: true,
+      citation: item.retriever_resources,
+      message_files: getProcessedFilesFromResponse(answerFiles.map((item: any) => ({ ...item, related_id: item.id }))),
+      parentMessageId: `question-${item.id}`,
+    })
+  })
+  return newChatList
+}
+
 export const useChatWithHistory = (installedAppInfo?: InstalledApp) => {
   const isInstalledApp = useMemo(() => !!installedAppInfo, [installedAppInfo])
   const { data: appInfo, isLoading: appInfoLoading, error: appInfoError } = useSWR(installedAppInfo ? null : 'appInfo', fetchAppInfo)
@@ -109,9 +138,9 @@ export const useChatWithHistory = (installedAppInfo?: InstalledApp) => {
   const { data: appConversationData, isLoading: appConversationDataLoading, mutate: mutateAppConversationData } = useSWR(['appConversationData', isInstalledApp, appId, false], () => fetchConversations(isInstalledApp, appId, undefined, false, 100))
   const { data: appChatListData, isLoading: appChatListDataLoading } = useSWR(chatShouldReloadKey ? ['appChatList', chatShouldReloadKey, isInstalledApp, appId] : null, () => fetchChatList(chatShouldReloadKey, isInstalledApp, appId))
 
-  const appPrevChatList = useMemo(
+  const appPrevChatTree = useMemo(
     () => (currentConversationId && appChatListData?.data.length)
-      ? getPrevChatList(appChatListData.data)
+      ? buildChatItemTree(getFormattedChatList(appChatListData.data))
       : [],
     [appChatListData, currentConversationId],
   )
@@ -403,7 +432,7 @@ export const useChatWithHistory = (installedAppInfo?: InstalledApp) => {
     appConversationDataLoading,
     appChatListData,
     appChatListDataLoading,
-    appPrevChatList,
+    appPrevChatTree,
     pinnedConversationList,
     conversationList,
     showConfigPanelBeforeChat,

web/app/components/base/chat/chat-with-history/index.tsx

@@ -30,7 +30,7 @@ const ChatWithHistory: FC<ChatWithHistoryProps> = ({
     appInfoError,
     appData,
     appInfoLoading,
-    appPrevChatList,
+    appPrevChatTree,
     showConfigPanelBeforeChat,
     appChatListDataLoading,
     chatShouldReloadKey,
@@ -38,7 +38,7 @@ const ChatWithHistory: FC<ChatWithHistoryProps> = ({
     themeBuilder,
   } = useChatWithHistoryContext()
 
-  const chatReady = (!showConfigPanelBeforeChat || !!appPrevChatList.length)
+  const chatReady = (!showConfigPanelBeforeChat || !!appPrevChatTree.length)
   const customConfig = appData?.custom_config
   const site = appData?.site
 
@@ -76,9 +76,9 @@ const ChatWithHistory: FC<ChatWithHistoryProps> = ({
           <HeaderInMobile />
         )
       }
-      <div className={`grow overflow-hidden ${showConfigPanelBeforeChat && !appPrevChatList.length && 'flex items-center justify-center'}`}>
+      <div className={`grow overflow-hidden ${showConfigPanelBeforeChat && !appPrevChatTree.length && 'flex items-center justify-center'}`}>
         {
-          showConfigPanelBeforeChat && !appChatListDataLoading && !appPrevChatList.length && (
+          showConfigPanelBeforeChat && !appChatListDataLoading && !appPrevChatTree.length && (
             <div className={`flex w-full items-center justify-center h-full ${isMobile && 'px-4'}`}>
               <ConfigPanel />
             </div>
@@ -120,7 +120,7 @@ const ChatWithHistoryWrap: FC<ChatWithHistoryWrapProps> = ({
     appChatListDataLoading,
     currentConversationId,
     currentConversationItem,
-    appPrevChatList,
+    appPrevChatTree,
     pinnedConversationList,
     conversationList,
     showConfigPanelBeforeChat,
@@ -154,7 +154,7 @@ const ChatWithHistoryWrap: FC<ChatWithHistoryWrapProps> = ({
       appChatListDataLoading,
       currentConversationId,
       currentConversationItem,
-      appPrevChatList,
+      appPrevChatTree,
       pinnedConversationList,
       conversationList,
       showConfigPanelBeforeChat,

web/app/components/base/chat/chat/answer/index.tsx

@@ -209,19 +209,19 @@ const Answer: FC<AnswerProps> = ({
             }
             {item.siblingCount && item.siblingCount > 1 && item.siblingIndex !== undefined && <div className="pt-3.5 flex justify-center items-center text-sm">
               <button
-                className={`${item.prevSibling ? 'opacity-100' : 'opacity-65'}`}
+                className={`${item.prevSibling ? 'opacity-100' : 'opacity-30'}`}
                 disabled={!item.prevSibling}
                 onClick={() => item.prevSibling && switchSibling?.(item.prevSibling)}
               >
-                <ChevronRight className="w-[14px] h-[14px] rotate-180 text-text-tertiary" />
+                <ChevronRight className="w-[14px] h-[14px] rotate-180 text-text-primary" />
               </button>
-              <span className="px-2 text-xs text-text-quaternary">{item.siblingIndex + 1} / {item.siblingCount}</span>
+              <span className="px-2 text-xs text-text-primary">{item.siblingIndex + 1} / {item.siblingCount}</span>
               <button
-                className={`${item.nextSibling ? 'opacity-100' : 'opacity-65'}`}
+                className={`${item.nextSibling ? 'opacity-100' : 'opacity-30'}`}
                 disabled={!item.nextSibling}
                 onClick={() => item.nextSibling && switchSibling?.(item.nextSibling)}
               >
-                <ChevronRight className="w-[14px] h-[14px] text-text-tertiary" />
+                <ChevronRight className="w-[14px] h-[14px] text-text-primary" />
               </button>
             </div>}
           </div>

web/app/components/base/chat/chat/hooks.ts

@@ -1,6 +1,7 @@
 import {
   useCallback,
   useEffect,
+  useMemo,
   useRef,
   useState,
 } from 'react'
@@ -12,8 +13,10 @@ import { v4 as uuidV4 } from 'uuid'
 import type {
   ChatConfig,
   ChatItem,
+  ChatItemInTree,
   Inputs,
 } from '../types'
+import { getThreadMessages } from '../utils'
 import type { InputForm } from './type'
 import {
   getProcessedInputs,
@@ -46,7 +49,7 @@ export const useChat = (
     inputs: Inputs
     inputsForm: InputForm[]
   },
-  prevChatList?: ChatItem[],
+  prevChatTree?: ChatItemInTree[],
   stopChat?: (taskId: string) => void,
 ) => {
   const { t } = useTranslation()
@@ -56,14 +59,48 @@ export const useChat = (
   const hasStopResponded = useRef(false)
   const [isResponding, setIsResponding] = useState(false)
   const isRespondingRef = useRef(false)
-  const [chatList, setChatList] = useState<ChatItem[]>(prevChatList || [])
-  const chatListRef = useRef<ChatItem[]>(prevChatList || [])
   const taskIdRef = useRef('')
   const [suggestedQuestions, setSuggestQuestions] = useState<string[]>([])
   const conversationMessagesAbortControllerRef = useRef<AbortController | null>(null)
   const suggestedQuestionsAbortControllerRef = useRef<AbortController | null>(null)
   const params = useParams()
   const pathname = usePathname()
+
+  const [chatTree, setChatTree] = useState<ChatItemInTree[]>(prevChatTree || [])
+  const chatTreeRef = useRef<ChatItemInTree[]>(chatTree)
+  const [targetMessageId, setTargetMessageId] = useState<string>()
+  const threadMessages = useMemo(() => getThreadMessages(chatTree, targetMessageId), [chatTree, targetMessageId])
+
+  const getIntroduction = useCallback((str: string) => {
+    return processOpeningStatement(str, formSettings?.inputs || {}, formSettings?.inputsForm || [])
+  }, [formSettings?.inputs, formSettings?.inputsForm])
+
+  /** Final chat list that will be rendered */
+  const chatList = useMemo(() => {
+    const ret = [...threadMessages]
+    if (config?.opening_statement) {
+      const index = threadMessages.findIndex(item => item.isOpeningStatement)
+
+      if (index > -1) {
+        ret[index] = {
+          ...ret[index],
+          content: getIntroduction(config.opening_statement),
+          suggestedQuestions: config.suggested_questions,
+        }
+      }
+      else {
+        ret.unshift({
+          id: `${Date.now()}`,
+          content: getIntroduction(config.opening_statement),
+          isAnswer: true,
+          isOpeningStatement: true,
+          suggestedQuestions: config.suggested_questions,
+        })
+      }
+    }
+    return ret
+  }, [threadMessages, config?.opening_statement, getIntroduction, config?.suggested_questions])
+
   useEffect(() => {
     setAutoFreeze(false)
     return () => {
@@ -71,43 +108,50 @@ export const useChat = (
     }
   }, [])
 
-  const handleUpdateChatList = useCallback((newChatList: ChatItem[]) => {
-    setChatList(newChatList)
-    chatListRef.current = newChatList
+  /** Find the target node by bfs and then operate on it */
+  const produceChatTreeNode = useCallback((targetId: string, operation: (node: ChatItemInTree) => void) => {
+    return produce(chatTreeRef.current, (draft) => {
+      const queue: ChatItemInTree[] = [...draft]
+      while (queue.length > 0) {
+        const current = queue.shift()!
+        if (current.id === targetId) {
+          operation(current)
+          break
+        }
+        if (current.children)
+          queue.push(...current.children)
+      }
+    })
   }, [])
+
+  type UpdateChatTreeNode = {
+    (id: string, fields: Partial<ChatItemInTree>): void
+    (id: string, update: (node: ChatItemInTree) => void): void
+  }
+
+  const updateChatTreeNode: UpdateChatTreeNode = useCallback((
+    id: string,
+    fieldsOrUpdate: Partial<ChatItemInTree> | ((node: ChatItemInTree) => void),
+  ) => {
+    const nextState = produceChatTreeNode(id, (node) => {
+      if (typeof fieldsOrUpdate === 'function') {
+        fieldsOrUpdate(node)
+      }
+      else {
+        Object.keys(fieldsOrUpdate).forEach((key) => {
+          (node as any)[key] = (fieldsOrUpdate as any)[key]
+        })
+      }
+    })
+    setChatTree(nextState)
+    chatTreeRef.current = nextState
+  }, [produceChatTreeNode])
+
   const handleResponding = useCallback((isResponding: boolean) => {
     setIsResponding(isResponding)
     isRespondingRef.current = isResponding
   }, [])
 
-  const getIntroduction = useCallback((str: string) => {
-    return processOpeningStatement(str, formSettings?.inputs || {}, formSettings?.inputsForm || [])
-  }, [formSettings?.inputs, formSettings?.inputsForm])
-  useEffect(() => {
-    if (config?.opening_statement) {
-      handleUpdateChatList(produce(chatListRef.current, (draft) => {
-        const index = draft.findIndex(item => item.isOpeningStatement)
-
-        if (index > -1) {
-          draft[index] = {
-            ...draft[index],
-            content: getIntroduction(config.opening_statement),
-            suggestedQuestions: config.suggested_questions,
-          }
-        }
-        else {
-          draft.unshift({
-            id: `${Date.now()}`,
-            content: getIntroduction(config.opening_statement),
-            isAnswer: true,
-            isOpeningStatement: true,
-            suggestedQuestions: config.suggested_questions,
-          })
-        }
-      }))
-    }
-  }, [config?.opening_statement, getIntroduction, config?.suggested_questions, handleUpdateChatList])
-
   const handleStop = useCallback(() => {
     hasStopResponded.current = true
     handleResponding(false)
@@ -123,50 +167,50 @@ export const useChat = (
     conversationId.current = ''
     taskIdRef.current = ''
     handleStop()
-    const newChatList = config?.opening_statement
-      ? [{
-        id: `${Date.now()}`,
-        content: config.opening_statement,
-        isAnswer: true,
-        isOpeningStatement: true,
-        suggestedQuestions: config.suggested_questions,
-      }]
-      : []
-    handleUpdateChatList(newChatList)
+    setChatTree([])
     setSuggestQuestions([])
-  }, [
-    config,
-    handleStop,
-    handleUpdateChatList,
-  ])
+  }, [handleStop])
 
-  const updateCurrentQA = useCallback(({
+  const updateCurrentQAOnTree = useCallback(({
+    parentId,
     responseItem,
-    questionId,
-    placeholderAnswerId,
+    placeholderQuestionId,
     questionItem,
   }: {
+    parentId?: string
     responseItem: ChatItem
-    questionId: string
-    placeholderAnswerId: string
+    placeholderQuestionId: 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 })
+    let nextState: ChatItemInTree[]
+    const currentQA = { ...questionItem, children: [{ ...responseItem, children: [] }] }
+    if (!parentId && !chatTree.some(item => [placeholderQuestionId, questionItem.id].includes(item.id))) {
+      // QA whose parent is not provided is considered as a first message of the conversation,
+      // and it should be a root node of the chat tree
+      nextState = produce(chatTree, (draft) => {
+        draft.push(currentQA)
+      })
+    }
+    else {
+      // find the target QA in the tree and update it; if not found, insert it to its parent node
+      nextState = produceChatTreeNode(parentId!, (parentNode) => {
+        const questionNodeIndex = parentNode.children!.findIndex(item => [placeholderQuestionId, questionItem.id].includes(item.id))
+        if (questionNodeIndex === -1)
+          parentNode.children!.push(currentQA)
+        else
+          parentNode.children![questionNodeIndex] = currentQA
       })
-    handleUpdateChatList(newListWithAnswer)
-  }, [handleUpdateChatList])
+    }
+    setChatTree(nextState)
+    chatTreeRef.current = nextState
+  }, [chatTree, produceChatTreeNode])
 
   const handleSend = useCallback(async (
     url: string,
     data: {
       query: string
       files?: FileEntity[]
+      parent_message_id?: string
       [key: string]: any
     },
     {
@@ -183,12 +227,15 @@ export const useChat = (
       return false
     }
 
-    const questionId = `question-${Date.now()}`
+    const parentMessage = threadMessages.find(item => item.id === data.parent_message_id)
+
+    const placeholderQuestionId = `question-${Date.now()}`
     const questionItem = {
-      id: questionId,
+      id: placeholderQuestionId,
       content: data.query,
       isAnswer: false,
       message_files: data.files,
+      parentMessageId: data.parent_message_id,
     }
 
     const placeholderAnswerId = `answer-placeholder-${Date.now()}`
@@ -196,18 +243,27 @@ export const useChat = (
       id: placeholderAnswerId,
       content: '',
       isAnswer: true,
+      parentMessageId: questionItem.id,
+      siblingIndex: parentMessage?.children?.length ?? chatTree.length,
     }
 
-    const newList = [...chatListRef.current, questionItem, placeholderAnswerItem]
-    handleUpdateChatList(newList)
+    setTargetMessageId(parentMessage?.id)
+    updateCurrentQAOnTree({
+      parentId: data.parent_message_id,
+      responseItem: placeholderAnswerItem,
+      placeholderQuestionId,
+      questionItem,
+    })
 
     // answer
-    const responseItem: ChatItem = {
+    const responseItem: ChatItemInTree = {
       id: placeholderAnswerId,
       content: '',
       agent_thoughts: [],
       message_files: [],
       isAnswer: true,
+      parentMessageId: questionItem.id,
+      siblingIndex: parentMessage?.children?.length ?? chatTree.length,
     }
 
     handleResponding(true)
@@ -268,7 +324,9 @@ export const useChat = (
           }
 
           if (messageId && !hasSetResponseId) {
+            questionItem.id = `question-${messageId}`
             responseItem.id = messageId
+            responseItem.parentMessageId = questionItem.id
             hasSetResponseId = true
           }
 
@@ -279,11 +337,11 @@ export const useChat = (
           if (messageId)
             responseItem.id = messageId
 
-          updateCurrentQA({
-            responseItem,
-            questionId,
-            placeholderAnswerId,
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
             questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
           })
         },
         async onCompleted(hasError?: boolean) {
@@ -304,43 +362,32 @@ export const useChat = (
             if (!newResponseItem)
               return
 
-            const newChatList = produce(chatListRef.current, (draft) => {
-              const index = draft.findIndex(item => item.id === responseItem.id)
-              if (index !== -1) {
-                const question = draft[index - 1]
-                draft[index - 1] = {
-                  ...question,
-                }
-                draft[index] = {
-                  ...draft[index],
-                  content: newResponseItem.answer,
-                  log: [
-                    ...newResponseItem.message,
-                    ...(newResponseItem.message[newResponseItem.message.length - 1].role !== 'assistant'
-                      ? [
-                        {
-                          role: 'assistant',
-                          text: newResponseItem.answer,
-                          files: newResponseItem.message_files?.filter((file: any) => file.belongs_to === 'assistant') || [],
-                        },
-                      ]
-                      : []),
-                  ],
-                  more: {
-                    time: formatTime(newResponseItem.created_at, 'hh:mm A'),
-                    tokens: newResponseItem.answer_tokens + newResponseItem.message_tokens,
-                    latency: newResponseItem.provider_response_latency.toFixed(2),
-                  },
-                  // for agent log
-                  conversationId: conversationId.current,
-                  input: {
-                    inputs: newResponseItem.inputs,
-                    query: newResponseItem.query,
-                  },
-                }
-              }
+            updateChatTreeNode(responseItem.id, {
+              content: newResponseItem.answer,
+              log: [
+                ...newResponseItem.message,
+                ...(newResponseItem.message[newResponseItem.message.length - 1].role !== 'assistant'
+                  ? [
+                    {
+                      role: 'assistant',
+                      text: newResponseItem.answer,
+                      files: newResponseItem.message_files?.filter((file: any) => file.belongs_to === 'assistant') || [],
+                    },
+                  ]
+                  : []),
+              ],
+              more: {
+                time: formatTime(newResponseItem.created_at, 'hh:mm A'),
+                tokens: newResponseItem.answer_tokens + newResponseItem.message_tokens,
+                latency: newResponseItem.provider_response_latency.toFixed(2),
+              },
+              // for agent log
+              conversationId: conversationId.current,
+              input: {
+                inputs: newResponseItem.inputs,
+                query: newResponseItem.query,
+              },
             })
-            handleUpdateChatList(newChatList)
           }
           if (config?.suggested_questions_after_answer?.enabled && !hasStopResponded.current && onGetSuggestedQuestions) {
             try {
@@ -360,11 +407,11 @@ export const useChat = (
           if (lastThought)
             responseItem.agent_thoughts![responseItem.agent_thoughts!.length - 1].message_files = [...(lastThought as any).message_files, file]
 
-          updateCurrentQA({
-            responseItem,
-            questionId,
-            placeholderAnswerId,
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
             questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
           })
         },
         onThought(thought) {
@@ -372,6 +419,7 @@ export const useChat = (
           const response = responseItem as any
           if (thought.message_id && !hasSetResponseId)
             response.id = thought.message_id
+
           if (response.agent_thoughts.length === 0) {
             response.agent_thoughts.push(thought)
           }
@@ -387,11 +435,11 @@ export const useChat = (
               responseItem.agent_thoughts!.push(thought)
             }
           }
-          updateCurrentQA({
-            responseItem,
-            questionId,
-            placeholderAnswerId,
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
             questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
           })
         },
         onMessageEnd: (messageEnd) => {
@@ -401,43 +449,36 @@ export const useChat = (
               id: messageEnd.metadata.annotation_reply.id,
               authorName: messageEnd.metadata.annotation_reply.account.name,
             })
-            const baseState = chatListRef.current.filter(item => item.id !== responseItem.id && item.id !== placeholderAnswerId)
-            const newListWithAnswer = produce(
-              baseState,
-              (draft) => {
-                if (!draft.find(item => item.id === questionId))
-                  draft.push({ ...questionItem })
-
-                draft.push({
-                  ...responseItem,
-                })
-              })
-            handleUpdateChatList(newListWithAnswer)
+            updateCurrentQAOnTree({
+              placeholderQuestionId,
+              questionItem,
+              responseItem,
+              parentId: data.parent_message_id,
+            })
             return
           }
           responseItem.citation = messageEnd.metadata?.retriever_resources || []
           const processedFilesFromResponse = getProcessedFilesFromResponse(messageEnd.files || [])
           responseItem.allFiles = uniqBy([...(responseItem.allFiles || []), ...(processedFilesFromResponse || [])], 'id')
 
-          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)
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
         },
         onMessageReplace: (messageReplace) => {
           responseItem.content = messageReplace.answer
         },
         onError() {
           handleResponding(false)
-          const newChatList = produce(chatListRef.current, (draft) => {
-            draft.splice(draft.findIndex(item => item.id === placeholderAnswerId), 1)
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
           })
-          handleUpdateChatList(newChatList)
         },
         onWorkflowStarted: ({ workflow_run_id, task_id }) => {
           taskIdRef.current = task_id
@@ -446,89 +487,84 @@ export const useChat = (
             status: WorkflowRunningStatus.Running,
             tracing: [],
           }
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
         },
-        onWorkflowFinished: ({ data }) => {
-          responseItem.workflowProcess!.status = data.status as WorkflowRunningStatus
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+        onWorkflowFinished: ({ data: workflowFinishedData }) => {
+          responseItem.workflowProcess!.status = workflowFinishedData.status as WorkflowRunningStatus
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
         },
-        onIterationStart: ({ data }) => {
+        onIterationStart: ({ data: iterationStartedData }) => {
           responseItem.workflowProcess!.tracing!.push({
-            ...data,
+            ...iterationStartedData,
             status: WorkflowRunningStatus.Running,
           } as any)
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
         },
-        onIterationFinish: ({ data }) => {
+        onIterationFinish: ({ data: iterationFinishedData }) => {
           const tracing = responseItem.workflowProcess!.tracing!
-          const iterationIndex = tracing.findIndex(item => item.node_id === data.node_id
-            && (item.execution_metadata?.parallel_id === data.execution_metadata?.parallel_id || item.parallel_id === data.execution_metadata?.parallel_id))!
+          const iterationIndex = tracing.findIndex(item => item.node_id === iterationFinishedData.node_id
+            && (item.execution_metadata?.parallel_id === iterationFinishedData.execution_metadata?.parallel_id || item.parallel_id === iterationFinishedData.execution_metadata?.parallel_id))!
           tracing[iterationIndex] = {
             ...tracing[iterationIndex],
-            ...data,
+            ...iterationFinishedData,
             status: WorkflowRunningStatus.Succeeded,
           } as any
 
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
         },
-        onNodeStarted: ({ data }) => {
-          if (data.iteration_id)
+        onNodeStarted: ({ data: nodeStartedData }) => {
+          if (nodeStartedData.iteration_id)
             return
 
           responseItem.workflowProcess!.tracing!.push({
-            ...data,
+            ...nodeStartedData,
             status: WorkflowRunningStatus.Running,
           } as any)
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
+          })
         },
-        onNodeFinished: ({ data }) => {
-          if (data.iteration_id)
+        onNodeFinished: ({ data: nodeFinishedData }) => {
+          if (nodeFinishedData.iteration_id)
             return
 
           const currentIndex = responseItem.workflowProcess!.tracing!.findIndex((item) => {
             if (!item.execution_metadata?.parallel_id)
-              return item.node_id === data.node_id
+              return item.node_id === nodeFinishedData.node_id
 
-            return item.node_id === data.node_id && (item.execution_metadata?.parallel_id === data.execution_metadata.parallel_id)
+            return item.node_id === nodeFinishedData.node_id && (item.execution_metadata?.parallel_id === nodeFinishedData.execution_metadata.parallel_id)
+          })
+          responseItem.workflowProcess!.tracing[currentIndex] = nodeFinishedData as any
+
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: data.parent_message_id,
           })
-          responseItem.workflowProcess!.tracing[currentIndex] = data as any
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
         },
         onTTSChunk: (messageId: string, audio: string) => {
           if (!audio || audio === '')
@@ -542,11 +578,13 @@ export const useChat = (
       })
     return true
   }, [
-    config?.suggested_questions_after_answer,
-    updateCurrentQA,
     t,
+    chatTree.length,
+    threadMessages,
+    config?.suggested_questions_after_answer,
+    updateCurrentQAOnTree,
+    updateChatTreeNode,
     notify,
-    handleUpdateChatList,
     handleResponding,
     formatTime,
     params.token,
@@ -556,76 +594,61 @@ export const useChat = (
   ])
 
   const handleAnnotationEdited = useCallback((query: string, answer: string, index: number) => {
-    handleUpdateChatList(chatListRef.current.map((item, i) => {
-      if (i === index - 1) {
-        return {
-          ...item,
-          content: query,
-        }
-      }
-      if (i === index) {
-        return {
-          ...item,
-          content: answer,
-          annotation: {
-            ...item.annotation,
-            logAnnotation: undefined,
-          } as any,
-        }
-      }
-      return item
-    }))
-  }, [handleUpdateChatList])
+    const targetQuestionId = chatList[index - 1].id
+    const targetAnswerId = chatList[index].id
+
+    updateChatTreeNode(targetQuestionId, {
+      content: query,
+    })
+    updateChatTreeNode(targetAnswerId, {
+      content: answer,
+      annotation: {
+        ...chatList[index].annotation,
+        logAnnotation: undefined,
+      } as any,
+    })
+  }, [chatList, updateChatTreeNode])
+
   const handleAnnotationAdded = useCallback((annotationId: string, authorName: string, query: string, answer: string, index: number) => {
-    handleUpdateChatList(chatListRef.current.map((item, i) => {
-      if (i === index - 1) {
-        return {
-          ...item,
-          content: query,
-        }
-      }
-      if (i === index) {
-        const answerItem = {
-          ...item,
-          content: item.content,
-          annotation: {
-            id: annotationId,
-            authorName,
-            logAnnotation: {
-              content: answer,
-              account: {
-                id: '',
-                name: authorName,
-                email: '',
-              },
-            },
-          } as Annotation,
-        }
-        return answerItem
-      }
-      return item
-    }))
-  }, [handleUpdateChatList])
-  const handleAnnotationRemoved = useCallback((index: number) => {
-    handleUpdateChatList(chatListRef.current.map((item, i) => {
-      if (i === index) {
-        return {
-          ...item,
-          content: item.content,
-          annotation: {
-            ...(item.annotation || {}),
+    const targetQuestionId = chatList[index - 1].id
+    const targetAnswerId = chatList[index].id
+
+    updateChatTreeNode(targetQuestionId, {
+      content: query,
+    })
+
+    updateChatTreeNode(targetAnswerId, {
+      content: chatList[index].content,
+      annotation: {
+        id: annotationId,
+        authorName,
+        logAnnotation: {
+          content: answer,
+          account: {
             id: '',
-          } as Annotation,
-        }
-      }
-      return item
-    }))
-  }, [handleUpdateChatList])
+            name: authorName,
+            email: '',
+          },
+        },
+      } as Annotation,
+    })
+  }, [chatList, updateChatTreeNode])
+
+  const handleAnnotationRemoved = useCallback((index: number) => {
+    const targetAnswerId = chatList[index].id
+
+    updateChatTreeNode(targetAnswerId, {
+      content: chatList[index].content,
+      annotation: {
+        ...(chatList[index].annotation || {}),
+        id: '',
+      } as Annotation,
+    })
+  }, [chatList, updateChatTreeNode])
 
   return {
     chatList,
-    chatListRef,
-    handleUpdateChatList,
+    setTargetMessageId,
     conversationId: conversationId.current,
     isResponding,
     setIsResponding,

web/app/components/base/chat/embedded-chatbot/chat-wrapper.tsx

@@ -3,10 +3,11 @@ import Chat from '../chat'
 import type {
   ChatConfig,
   ChatItem,
+  ChatItemInTree,
   OnSend,
 } from '../types'
 import { useChat } from '../chat/hooks'
-import { getLastAnswer } from '../utils'
+import { getLastAnswer, isValidGeneratedAnswer } from '../utils'
 import { useEmbeddedChatbotContext } from './context'
 import ConfigPanel from './config-panel'
 import { isDify } from './utils'
@@ -51,13 +52,12 @@ const ChatWrapper = () => {
     } as ChatConfig
   }, [appParams, currentConversationItem?.introduction, currentConversationId])
   const {
-    chatListRef,
     chatList,
+    setTargetMessageId,
     handleSend,
     handleStop,
     isResponding,
     suggestedQuestions,
-    handleUpdateChatList,
   } = useChat(
     appConfig,
     {
@@ -71,15 +71,15 @@ const ChatWrapper = () => {
   useEffect(() => {
     if (currentChatInstanceRef.current)
       currentChatInstanceRef.current.handleStop = handleStop
-  }, [])
+  }, [currentChatInstanceRef, handleStop])
 
-  const doSend: OnSend = useCallback((message, files, last_answer) => {
+  const doSend: OnSend = useCallback((message, files, isRegenerate = false, parentAnswer: ChatItem | null = null) => {
     const data: any = {
       query: message,
       files,
       inputs: currentConversationId ? currentConversationItem?.inputs : newConversationInputs,
       conversation_id: currentConversationId,
-      parent_message_id: last_answer?.id || getLastAnswer(chatListRef.current)?.id || null,
+      parent_message_id: (isRegenerate ? parentAnswer?.id : getLastAnswer(chatList)?.id) || null,
     }
 
     handleSend(
@@ -92,32 +92,21 @@ const ChatWrapper = () => {
       },
     )
   }, [
-    chatListRef,
-    appConfig,
+    chatList,
+    handleNewConversationCompleted,
+    handleSend,
     currentConversationId,
     currentConversationItem,
-    handleSend,
     newConversationInputs,
-    handleNewConversationCompleted,
     isInstalledApp,
     appId,
   ])
 
-  const doRegenerate = useCallback((chatItem: ChatItem) => {
-    const index = chatList.findIndex(item => item.id === chatItem.id)
-    if (index === -1)
-      return
-
-    const prevMessages = chatList.slice(0, index)
-    const question = prevMessages.pop()
-    const lastAnswer = getLastAnswer(prevMessages)
-
-    if (!question)
-      return
-
-    handleUpdateChatList(prevMessages)
-    doSend(question.content, question.message_files, lastAnswer)
-  }, [chatList, handleUpdateChatList, doSend])
+  const doRegenerate = useCallback((chatItem: ChatItemInTree) => {
+    const question = chatList.find(item => item.id === chatItem.parentMessageId)!
+    const parentAnswer = chatList.find(item => item.id === question.parentMessageId)
+    doSend(question.content, question.message_files, true, isValidGeneratedAnswer(parentAnswer) ? parentAnswer : null)
+  }, [chatList, doSend])
 
   const chatNode = useMemo(() => {
     if (inputsForms.length) {
@@ -172,6 +161,7 @@ const ChatWrapper = () => {
       answerIcon={answerIcon}
       hideProcessDetail
       themeBuilder={themeBuilder}
+      switchSibling={siblingMessageId => setTargetMessageId(siblingMessageId)}
     />
   )
 }

web/app/components/base/chat/types.ts

@@ -67,9 +67,12 @@ export type ChatItem = IChatItem & {
 
 export type ChatItemInTree = {
   children?: ChatItemInTree[]
-} & IChatItem
+} & ChatItem
 
-export type OnSend = (message: string, files?: FileEntity[], last_answer?: ChatItem | null) => void
+export type OnSend = {
+  (message: string, files?: FileEntity[]): void
+  (message: string, files: FileEntity[] | undefined, isRegenerate: boolean, lastAnswer?: ChatItem | null): void
+}
 
 export type OnRegenerate = (chatItem: ChatItem) => void
 

web/app/components/base/chat/utils.ts

@@ -1,8 +1,6 @@
-import { addFileInfos, sortAgentSorts } from '../../tools/utils'
 import { UUID_NIL } from './constants'
 import type { IChatItem } from './chat/type'
 import type { ChatItem, ChatItemInTree } from './types'
-import { getProcessedFilesFromResponse } from '@/app/components/base/file-uploader/utils'
 
 async function decodeBase64AndDecompress(base64String: string) {
   const binaryString = atob(base64String)
@@ -21,67 +19,24 @@ function getProcessedInputsFromUrlParams(): Record<string, any> {
   return inputs
 }
 
-function getLastAnswer(chatList: ChatItem[]) {
+function isValidGeneratedAnswer(item?: ChatItem | ChatItemInTree): boolean {
+  return !!item && item.isAnswer && !item.id.startsWith('answer-placeholder-') && !item.isOpeningStatement
+}
+
+function getLastAnswer<T extends ChatItem | ChatItemInTree>(chatList: T[]): T | null {
   for (let i = chatList.length - 1; i >= 0; i--) {
     const item = chatList[i]
-    if (item.isAnswer && !item.id.startsWith('answer-placeholder-') && !item.isOpeningStatement)
+    if (isValidGeneratedAnswer(item))
       return item
   }
   return null
 }
 
-function appendQAToChatList(chatList: ChatItem[], item: any) {
-  // we append answer first and then question since will reverse the whole chatList later
-  const answerFiles = item.message_files?.filter((file: any) => file.belongs_to === 'assistant') || []
-  chatList.push({
-    id: item.id,
-    content: item.answer,
-    agent_thoughts: addFileInfos(item.agent_thoughts ? sortAgentSorts(item.agent_thoughts) : item.agent_thoughts, item.message_files),
-    feedback: item.feedback,
-    isAnswer: true,
-    citation: item.retriever_resources,
-    message_files: getProcessedFilesFromResponse(answerFiles.map((item: any) => ({ ...item, related_id: item.id }))),
-  })
-  const questionFiles = item.message_files?.filter((file: any) => file.belongs_to === 'user') || []
-  chatList.push({
-    id: `question-${item.id}`,
-    content: item.query,
-    isAnswer: false,
-    message_files: getProcessedFilesFromResponse(questionFiles.map((item: any) => ({ ...item, related_id: item.id }))),
-  })
-}
-
 /**
- * Computes the latest thread messages from all messages of the conversation.
- * Same logic as backend codebase `api/core/prompt/utils/extract_thread_messages.py`
- *
- * @param fetchedMessages - The history chat list data from the backend, sorted by created_at in descending order. This includes all flattened history messages of the conversation.
- * @returns An array of ChatItems representing the latest thread.
+ * Build a chat item tree from a chat list
+ * @param allMessages - The chat list, sorted from oldest to newest
+ * @returns The chat item tree
  */
-function getPrevChatList(fetchedMessages: any[]) {
-  const ret: ChatItem[] = []
-  let nextMessageId = null
-
-  for (const item of fetchedMessages) {
-    if (!item.parent_message_id) {
-      appendQAToChatList(ret, item)
-      break
-    }
-
-    if (!nextMessageId) {
-      appendQAToChatList(ret, item)
-      nextMessageId = item.parent_message_id
-    }
-    else {
-      if (item.id === nextMessageId || nextMessageId === UUID_NIL) {
-        appendQAToChatList(ret, item)
-        nextMessageId = item.parent_message_id
-      }
-    }
-  }
-  return ret.reverse()
-}
-
 function buildChatItemTree(allMessages: IChatItem[]): ChatItemInTree[] {
   const map: Record<string, ChatItemInTree> = {}
   const rootNodes: ChatItemInTree[] = []
@@ -208,7 +163,7 @@ function getThreadMessages(tree: ChatItemInTree[], targetMessageId?: string): Ch
 
 export {
   getProcessedInputsFromUrlParams,
-  getPrevChatList,
+  isValidGeneratedAnswer,
   getLastAnswer,
   buildChatItemTree,
   getThreadMessages,

web/app/components/workflow/panel/debug-and-preview/chat-wrapper.tsx

@@ -19,14 +19,14 @@ import ConversationVariableModal from './conversation-variable-modal'
 import { useChat } from './hooks'
 import type { ChatWrapperRefType } from './index'
 import Chat from '@/app/components/base/chat/chat'
-import type { ChatItem, OnSend } from '@/app/components/base/chat/types'
+import type { ChatItem, ChatItemInTree, OnSend } from '@/app/components/base/chat/types'
 import { useFeatures } from '@/app/components/base/features/hooks'
 import {
   fetchSuggestedQuestions,
   stopChatMessageResponding,
 } from '@/service/debug'
 import { useStore as useAppStore } from '@/app/components/app/store'
-import { getLastAnswer } from '@/app/components/base/chat/utils'
+import { getLastAnswer, isValidGeneratedAnswer } from '@/app/components/base/chat/utils'
 
 type ChatWrapperProps = {
   showConversationVariableModal: boolean
@@ -65,13 +65,12 @@ const ChatWrapper = forwardRef<ChatWrapperRefType, ChatWrapperProps>(({
   const {
     conversationId,
     chatList,
-    chatListRef,
-    handleUpdateChatList,
     handleStop,
     isResponding,
     suggestedQuestions,
     handleSend,
     handleRestart,
+    setTargetMessageId,
   } = useChat(
     config,
     {
@@ -82,36 +81,26 @@ const ChatWrapper = forwardRef<ChatWrapperRefType, ChatWrapperProps>(({
     taskId => stopChatMessageResponding(appDetail!.id, taskId),
   )
 
-  const doSend = useCallback<OnSend>((query, files, last_answer) => {
+  const doSend: OnSend = useCallback((message, files, isRegenerate = false, parentAnswer: ChatItem | null = null) => {
     handleSend(
       {
-        query,
+        query: message,
         files,
         inputs: workflowStore.getState().inputs,
         conversation_id: conversationId,
-        parent_message_id: last_answer?.id || getLastAnswer(chatListRef.current)?.id || null,
+        parent_message_id: (isRegenerate ? parentAnswer?.id : getLastAnswer(chatList)?.id) || undefined,
       },
       {
         onGetSuggestedQuestions: (messageId, getAbortController) => fetchSuggestedQuestions(appDetail!.id, messageId, getAbortController),
       },
     )
-  }, [chatListRef, conversationId, handleSend, workflowStore, appDetail])
+  }, [handleSend, workflowStore, conversationId, chatList, appDetail])
 
-  const doRegenerate = useCallback((chatItem: ChatItem) => {
-    const index = chatList.findIndex(item => item.id === chatItem.id)
-    if (index === -1)
-      return
-
-    const prevMessages = chatList.slice(0, index)
-    const question = prevMessages.pop()
-    const lastAnswer = getLastAnswer(prevMessages)
-
-    if (!question)
-      return
-
-    handleUpdateChatList(prevMessages)
-    doSend(question.content, question.message_files, lastAnswer)
-  }, [chatList, handleUpdateChatList, doSend])
+  const doRegenerate = useCallback((chatItem: ChatItemInTree) => {
+    const question = chatList.find(item => item.id === chatItem.parentMessageId)!
+    const parentAnswer = chatList.find(item => item.id === question.parentMessageId)
+    doSend(question.content, question.message_files, true, isValidGeneratedAnswer(parentAnswer) ? parentAnswer : null)
+  }, [chatList, doSend])
 
   useImperativeHandle(ref, () => {
     return {
@@ -159,6 +148,7 @@ const ChatWrapper = forwardRef<ChatWrapperRefType, ChatWrapperProps>(({
         suggestedQuestions={suggestedQuestions}
         showPromptLog
         chatAnswerContainerInner='!pr-2'
+        switchSibling={setTargetMessageId}
       />
       {showConversationVariableModal && (
         <ConversationVariableModal

web/app/components/workflow/panel/debug-and-preview/hooks.ts

@@ -1,6 +1,7 @@
 import {
   useCallback,
   useEffect,
+  useMemo,
   useRef,
   useState,
 } from 'react'
@@ -13,6 +14,7 @@ import { useWorkflowStore } from '../../store'
 import { DEFAULT_ITER_TIMES } from '../../constants'
 import type {
   ChatItem,
+  ChatItemInTree,
   Inputs,
 } from '@/app/components/base/chat/types'
 import type { InputForm } from '@/app/components/base/chat/chat/type'
@@ -27,6 +29,7 @@ import {
   getProcessedFilesFromResponse,
 } from '@/app/components/base/file-uploader/utils'
 import type { FileEntity } from '@/app/components/base/file-uploader/types'
+import { getThreadMessages } from '@/app/components/base/chat/utils'
 import type { NodeTracing } from '@/types/workflow'
 
 type GetAbortController = (abortController: AbortController) => void
@@ -39,7 +42,7 @@ export const useChat = (
     inputs: Inputs
     inputsForm: InputForm[]
   },
-  prevChatList?: ChatItem[],
+  prevChatTree?: ChatItemInTree[],
   stopChat?: (taskId: string) => void,
 ) => {
   const { t } = useTranslation()
@@ -49,60 +52,76 @@ export const useChat = (
   const workflowStore = useWorkflowStore()
   const conversationId = useRef('')
   const taskIdRef = useRef('')
-  const [chatList, setChatList] = useState<ChatItem[]>(prevChatList || [])
-  const chatListRef = useRef<ChatItem[]>(prevChatList || [])
   const [isResponding, setIsResponding] = useState(false)
   const isRespondingRef = useRef(false)
   const [suggestedQuestions, setSuggestQuestions] = useState<string[]>([])
   const suggestedQuestionsAbortControllerRef = useRef<AbortController | null>(null)
-
   const {
     setIterTimes,
   } = workflowStore.getState()
-  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 [chatTree, setChatTree] = useState<ChatItemInTree[]>(prevChatTree || [])
+  const chatTreeRef = useRef<ChatItemInTree[]>(chatTree)
+  const [targetMessageId, setTargetMessageId] = useState<string>()
+  const threadMessages = useMemo(() => getThreadMessages(chatTree, targetMessageId), [chatTree, targetMessageId])
+
   const getIntroduction = useCallback((str: string) => {
     return processOpeningStatement(str, formSettings?.inputs || {}, formSettings?.inputsForm || [])
   }, [formSettings?.inputs, formSettings?.inputsForm])
-  useEffect(() => {
+
+  /** Final chat list that will be rendered */
+  const chatList = useMemo(() => {
+    const ret = [...threadMessages]
     if (config?.opening_statement) {
-      handleUpdateChatList(produce(chatListRef.current, (draft) => {
-        const index = draft.findIndex(item => item.isOpeningStatement)
-
-        if (index > -1) {
-          draft[index] = {
-            ...draft[index],
-            content: getIntroduction(config.opening_statement),
-            suggestedQuestions: config.suggested_questions,
-          }
-        }
-        else {
-          draft.unshift({
-            id: `${Date.now()}`,
-            content: getIntroduction(config.opening_statement),
-            isAnswer: true,
-            isOpeningStatement: true,
-            suggestedQuestions: config.suggested_questions,
-          })
+      const index = threadMessages.findIndex(item => item.isOpeningStatement)
+
+      if (index > -1) {
+        ret[index] = {
+          ...ret[index],
+          content: getIntroduction(config.opening_statement),
+          suggestedQuestions: config.suggested_questions,
         }
-      }))
+      }
+      else {
+        ret.unshift({
+          id: `${Date.now()}`,
+          content: getIntroduction(config.opening_statement),
+          isAnswer: true,
+          isOpeningStatement: true,
+          suggestedQuestions: config.suggested_questions,
+        })
+      }
+    }
+    return ret
+  }, [threadMessages, config?.opening_statement, getIntroduction, config?.suggested_questions])
+
+  useEffect(() => {
+    setAutoFreeze(false)
+    return () => {
+      setAutoFreeze(true)
     }
-  }, [config?.opening_statement, getIntroduction, config?.suggested_questions, handleUpdateChatList])
+  }, [])
+
+  /** Find the target node by bfs and then operate on it */
+  const produceChatTreeNode = useCallback((targetId: string, operation: (node: ChatItemInTree) => void) => {
+    return produce(chatTreeRef.current, (draft) => {
+      const queue: ChatItemInTree[] = [...draft]
+      while (queue.length > 0) {
+        const current = queue.shift()!
+        if (current.id === targetId) {
+          operation(current)
+          break
+        }
+        if (current.children)
+          queue.push(...current.children)
+      }
+    })
+  }, [])
 
   const handleStop = useCallback(() => {
     hasStopResponded.current = true
@@ -119,50 +138,52 @@ export const useChat = (
     taskIdRef.current = ''
     handleStop()
     setIterTimes(DEFAULT_ITER_TIMES)
-    const newChatList = config?.opening_statement
-      ? [{
-        id: `${Date.now()}`,
-        content: config.opening_statement,
-        isAnswer: true,
-        isOpeningStatement: true,
-        suggestedQuestions: config.suggested_questions,
-      }]
-      : []
-    handleUpdateChatList(newChatList)
+    setChatTree([])
     setSuggestQuestions([])
   }, [
-    config,
     handleStop,
-    handleUpdateChatList,
     setIterTimes,
   ])
 
-  const updateCurrentQA = useCallback(({
+  const updateCurrentQAOnTree = useCallback(({
+    parentId,
     responseItem,
-    questionId,
-    placeholderAnswerId,
+    placeholderQuestionId,
     questionItem,
   }: {
+    parentId?: string
     responseItem: ChatItem
-    questionId: string
-    placeholderAnswerId: string
+    placeholderQuestionId: 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 })
+    let nextState: ChatItemInTree[]
+    const currentQA = { ...questionItem, children: [{ ...responseItem, children: [] }] }
+    if (!parentId && !chatTree.some(item => [placeholderQuestionId, questionItem.id].includes(item.id))) {
+      // QA whose parent is not provided is considered as a first message of the conversation,
+      // and it should be a root node of the chat tree
+      nextState = produce(chatTree, (draft) => {
+        draft.push(currentQA)
+      })
+    }
+    else {
+      // find the target QA in the tree and update it; if not found, insert it to its parent node
+      nextState = produceChatTreeNode(parentId!, (parentNode) => {
+        const questionNodeIndex = parentNode.children!.findIndex(item => [placeholderQuestionId, questionItem.id].includes(item.id))
+        if (questionNodeIndex === -1)
+          parentNode.children!.push(currentQA)
+        else
+          parentNode.children![questionNodeIndex] = currentQA
       })
-    handleUpdateChatList(newListWithAnswer)
-  }, [handleUpdateChatList])
+    }
+    setChatTree(nextState)
+    chatTreeRef.current = nextState
+  }, [chatTree, produceChatTreeNode])
 
   const handleSend = useCallback((
     params: {
       query: string
       files?: FileEntity[]
+      parent_message_id?: string
       [key: string]: any
     },
     {
@@ -174,12 +195,15 @@ export const useChat = (
       return false
     }
 
-    const questionId = `question-${Date.now()}`
+    const parentMessage = threadMessages.find(item => item.id === params.parent_message_id)
+
+    const placeholderQuestionId = `question-${Date.now()}`
     const questionItem = {
-      id: questionId,
+      id: placeholderQuestionId,
       content: params.query,
       isAnswer: false,
       message_files: params.files,
+      parentMessageId: params.parent_message_id,
     }
 
     const placeholderAnswerId = `answer-placeholder-${Date.now()}`
@@ -187,10 +211,17 @@ export const useChat = (
       id: placeholderAnswerId,
       content: '',
       isAnswer: true,
+      parentMessageId: questionItem.id,
+      siblingIndex: parentMessage?.children?.length ?? chatTree.length,
     }
 
-    const newList = [...chatListRef.current, questionItem, placeholderAnswerItem]
-    handleUpdateChatList(newList)
+    setTargetMessageId(parentMessage?.id)
+    updateCurrentQAOnTree({
+      parentId: params.parent_message_id,
+      responseItem: placeholderAnswerItem,
+      placeholderQuestionId,
+      questionItem,
+    })
 
     // answer
     const responseItem: ChatItem = {
@@ -199,6 +230,8 @@ export const useChat = (
       agent_thoughts: [],
       message_files: [],
       isAnswer: true,
+      parentMessageId: questionItem.id,
+      siblingIndex: parentMessage?.children?.length ?? chatTree.length,
     }
 
     handleResponding(true)
@@ -230,7 +263,9 @@ export const useChat = (
           responseItem.content = responseItem.content + message
 
           if (messageId && !hasSetResponseId) {
+            questionItem.id = `question-${messageId}`
             responseItem.id = messageId
+            responseItem.parentMessageId = questionItem.id
             hasSetResponseId = true
           }
 
@@ -241,11 +276,11 @@ export const useChat = (
           if (messageId)
             responseItem.id = messageId
 
-          updateCurrentQA({
-            responseItem,
-            questionId,
-            placeholderAnswerId,
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
             questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
           })
         },
         async onCompleted(hasError?: boolean, errorMessage?: string) {
@@ -255,15 +290,12 @@ export const useChat = (
             if (errorMessage) {
               responseItem.content = errorMessage
               responseItem.isError = true
-              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)
+              updateCurrentQAOnTree({
+                placeholderQuestionId,
+                questionItem,
+                responseItem,
+                parentId: params.parent_message_id,
+              })
             }
             return
           }
@@ -286,15 +318,12 @@ export const useChat = (
           const processedFilesFromResponse = getProcessedFilesFromResponse(messageEnd.files || [])
           responseItem.allFiles = uniqBy([...(responseItem.allFiles || []), ...(processedFilesFromResponse || [])], 'id')
 
-          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)
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onMessageReplace: (messageReplace) => {
           responseItem.content = messageReplace.answer
@@ -309,23 +338,21 @@ export const useChat = (
             status: WorkflowRunningStatus.Running,
             tracing: [],
           }
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onWorkflowFinished: ({ data }) => {
           responseItem.workflowProcess!.status = data.status as WorkflowRunningStatus
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onIterationStart: ({ data }) => {
           responseItem.workflowProcess!.tracing!.push({
@@ -333,13 +360,12 @@ export const useChat = (
             status: NodeRunningStatus.Running,
             details: [],
           } as any)
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onIterationNext: ({ data }) => {
           const tracing = responseItem.workflowProcess!.tracing!
@@ -347,10 +373,12 @@ export const useChat = (
             && (item.execution_metadata?.parallel_id === data.execution_metadata?.parallel_id || item.parallel_id === data.execution_metadata?.parallel_id))!
           iterations.details!.push([])
 
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.length - 1
-            draft[currentIndex] = responseItem
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onIterationFinish: ({ data }) => {
           const tracing = responseItem.workflowProcess!.tracing!
@@ -361,10 +389,12 @@ export const useChat = (
             ...data,
             status: NodeRunningStatus.Succeeded,
           } as any
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.length - 1
-            draft[currentIndex] = responseItem
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onNodeStarted: ({ data }) => {
           if (data.iteration_id)
@@ -374,13 +404,12 @@ export const useChat = (
             ...data,
             status: NodeRunningStatus.Running,
           } as any)
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
         onNodeRetry: ({ data }) => {
           if (data.iteration_id)
@@ -422,23 +451,21 @@ export const useChat = (
               : {}),
             ...data,
           } as any
-          handleUpdateChatList(produce(chatListRef.current, (draft) => {
-            const currentIndex = draft.findIndex(item => item.id === responseItem.id)
-            draft[currentIndex] = {
-              ...draft[currentIndex],
-              ...responseItem,
-            }
-          }))
+          updateCurrentQAOnTree({
+            placeholderQuestionId,
+            questionItem,
+            responseItem,
+            parentId: params.parent_message_id,
+          })
         },
       },
     )
-  }, [handleRun, handleResponding, handleUpdateChatList, notify, t, updateCurrentQA, config.suggested_questions_after_answer?.enabled, formSettings])
+  }, [threadMessages, chatTree.length, updateCurrentQAOnTree, handleResponding, formSettings?.inputsForm, handleRun, notify, t, config?.suggested_questions_after_answer?.enabled])
 
   return {
     conversationId: conversationId.current,
     chatList,
-    chatListRef,
-    handleUpdateChatList,
+    setTargetMessageId,
     handleSend,
     handleStop,
     handleRestart,