feat: Support for conversational streaming (#809)

### What problem does this PR solve? feat: Support for conversational streaming #709 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
2025-12-08 20:42:30 +08:00 · 2024-05-16 20:15:02 +08:00
parent 95f809187e
commit c6c9dbde64
21 changed files with 424 additions and 255 deletions
--- a/web/src/pages/chat/chat-container/index.tsx
+++ b/web/src/pages/chat/chat-container/index.tsx
@ -6,16 +6,7 @@ import { useSelectFileThumbnails } from '@/hooks/knowledgeHook';
 import { useSelectUserInfo } from '@/hooks/userSettingHook';
 import { IReference, Message } from '@/interfaces/database/chat';
 import { IChunk } from '@/interfaces/database/knowledge';
-import {
-  Avatar,
-  Button,
-  Drawer,
-  Flex,
-  Input,
-  List,
-  Skeleton,
-  Spin,
-} from 'antd';
+import { Avatar, Button, Drawer, Flex, Input, List, Spin } from 'antd';
 import classNames from 'classnames';
 import { useMemo } from 'react';
 import {
@ -32,20 +23,24 @@ import SvgIcon from '@/components/svg-icon';
 import { useTranslate } from '@/hooks/commonHooks';
 import { useGetDocumentUrl } from '@/hooks/documentHooks';
 import { getExtension, isPdf } from '@/utils/documentUtils';
+import { buildMessageItemReference } from '../utils';
 import styles from './index.less';

 const MessageItem = ({
  item,
  reference,
+  loading = false,
  clickDocumentButton,
 }: {
  item: Message;
  reference: IReference;
+  loading?: boolean;
  clickDocumentButton: (documentId: string, chunk: IChunk) => void;
 }) => {
  const userInfo = useSelectUserInfo();
  const fileThumbnails = useSelectFileThumbnails();
  const getDocumentUrl = useGetDocumentUrl();
+  const { t } = useTranslate('chat');

  const isAssistant = item.role === MessageType.Assistant;

@ -53,6 +48,14 @@ const MessageItem = ({
    return reference?.doc_aggs ?? [];
  }, [reference?.doc_aggs]);

+  const content = useMemo(() => {
+    let text = item.content;
+    if (text === '') {
+      text = t('searching');
+    }
+    return loading ? text?.concat('~~2$$') : text;
+  }, [item.content, loading, t]);
+
  return (
    <div
      className={classNames(styles.messageItem, {
@ -85,15 +88,11 @@ const MessageItem = ({
          <Flex vertical gap={8} flex={1}>
            <b>{isAssistant ? '' : userInfo.nickname}</b>
            <div className={styles.messageText}>
-              {item.content !== '' ? (
-                <MarkdownContent
-                  content={item.content}
-                  reference={reference}
-                  clickDocumentButton={clickDocumentButton}
-                ></MarkdownContent>
-              ) : (
-                <Skeleton active className={styles.messageEmpty} />
-              )}
+              <MarkdownContent
+                content={content}
+                reference={reference}
+                clickDocumentButton={clickDocumentButton}
+              ></MarkdownContent>
            </div>
            {isAssistant && referenceDocumentList.length > 0 && (
              <List
@ -139,13 +138,19 @@ const ChatContainer = () => {
    currentConversation: conversation,
    addNewestConversation,
    removeLatestMessage,
+    addNewestAnswer,
  } = useFetchConversationOnMount();
  const {
    handleInputChange,
    handlePressEnter,
    value,
    loading: sendLoading,
-  } = useSendMessage(conversation, addNewestConversation, removeLatestMessage);
+  } = useSendMessage(
+    conversation,
+    addNewestConversation,
+    removeLatestMessage,
+    addNewestAnswer,
+  );
  const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
    useClickDrawer();
  const disabled = useGetSendButtonDisabled();
@ -159,19 +164,17 @@ const ChatContainer = () => {
        <Flex flex={1} vertical className={styles.messageContainer}>
          <div>
            <Spin spinning={loading}>
-              {conversation?.message?.map((message) => {
-                const assistantMessages = conversation?.message
-                  ?.filter((x) => x.role === MessageType.Assistant)
-                  .slice(1);
-                const referenceIndex = assistantMessages.findIndex(
-                  (x) => x.id === message.id,
-                );
-                const reference = conversation.reference[referenceIndex];
+              {conversation?.message?.map((message, i) => {
                return (
                  <MessageItem
+                    loading={
+                      message.role === MessageType.Assistant &&
+                      sendLoading &&
+                      conversation?.message.length - 1 === i
+                    }
                    key={message.id}
                    item={message}
-                    reference={reference}
+                    reference={buildMessageItemReference(conversation, message)}
                    clickDocumentButton={clickDocumentButton}
                  ></MessageItem>
                );
--- a/web/src/pages/chat/hooks.ts
+++ b/web/src/pages/chat/hooks.ts
@ -1,7 +1,6 @@
 import { MessageType } from '@/constants/chat';
 import { fileIconMap } from '@/constants/common';
 import {
-  useCompleteConversation,
  useCreateToken,
  useFetchConversation,
  useFetchConversationList,
@ -24,8 +23,14 @@ import {
  useShowDeleteConfirm,
  useTranslate,
 } from '@/hooks/commonHooks';
+import { useSendMessageWithSse } from '@/hooks/logicHooks';
 import { useOneNamespaceEffectsLoading } from '@/hooks/storeHooks';
-import { IConversation, IDialog, IStats } from '@/interfaces/database/chat';
+import {
+  IAnswer,
+  IConversation,
+  IDialog,
+  IStats,
+} from '@/interfaces/database/chat';
 import { IChunk } from '@/interfaces/database/knowledge';
 import { getFileExtension } from '@/utils';
 import { message } from 'antd';
@ -380,31 +385,56 @@ export const useSelectCurrentConversation = () => {
  const dialog = useSelectCurrentDialog();
  const { conversationId, dialogId } = useGetChatSearchParams();

-  const addNewestConversation = useCallback((message: string) => {
+  const addNewestConversation = useCallback(
+    (message: string, answer: string = '') => {
+      setCurrentConversation((pre) => {
+        return {
+          ...pre,
+          message: [
+            ...pre.message,
+            {
+              role: MessageType.User,
+              content: message,
+              id: uuid(),
+            } as IMessage,
+            {
+              role: MessageType.Assistant,
+              content: answer,
+              id: uuid(),
+              reference: [],
+            } as IMessage,
+          ],
+        };
+      });
+    },
+    [],
+  );
+
+  const addNewestAnswer = useCallback((answer: IAnswer) => {
    setCurrentConversation((pre) => {
-      return {
-        ...pre,
-        message: [
-          ...pre.message,
-          {
-            role: MessageType.User,
-            content: message,
-            id: uuid(),
-          } as IMessage,
-          {
-            role: MessageType.Assistant,
-            content: '',
-            id: uuid(),
-            reference: [],
-          } as IMessage,
-        ],
-      };
+      const latestMessage = pre.message?.at(-1);
+
+      if (latestMessage) {
+        return {
+          ...pre,
+          message: [
+            ...pre.message.slice(0, -1),
+            {
+              ...latestMessage,
+              content: answer.answer,
+              reference: answer.reference,
+            } as IMessage,
+          ],
+        };
+      }
+      return pre;
    });
  }, []);

  const removeLatestMessage = useCallback(() => {
+    console.info('removeLatestMessage');
    setCurrentConversation((pre) => {
-      const nextMessages = pre.message.slice(0, -2);
+      const nextMessages = pre.message?.slice(0, -2) ?? [];
      return {
        ...pre,
        message: nextMessages,
@ -441,7 +471,12 @@ export const useSelectCurrentConversation = () => {
    }
  }, [conversation, conversationId]);

-  return { currentConversation, addNewestConversation, removeLatestMessage };
+  return {
+    currentConversation,
+    addNewestConversation,
+    removeLatestMessage,
+    addNewestAnswer,
+  };
 };

 export const useScrollToBottom = (currentConversation: IClientConversation) => {
@ -464,8 +499,12 @@ export const useScrollToBottom = (currentConversation: IClientConversation) => {
 export const useFetchConversationOnMount = () => {
  const { conversationId } = useGetChatSearchParams();
  const fetchConversation = useFetchConversation();
-  const { currentConversation, addNewestConversation, removeLatestMessage } =
-    useSelectCurrentConversation();
+  const {
+    currentConversation,
+    addNewestConversation,
+    removeLatestMessage,
+    addNewestAnswer,
+  } = useSelectCurrentConversation();
  const ref = useScrollToBottom(currentConversation);

  const fetchConversationOnMount = useCallback(() => {
@ -483,6 +522,7 @@ export const useFetchConversationOnMount = () => {
    addNewestConversation,
    ref,
    removeLatestMessage,
+    addNewestAnswer,
  };
 };

@ -504,25 +544,22 @@ export const useHandleMessageInputChange = () => {

 export const useSendMessage = (
  conversation: IClientConversation,
-  addNewestConversation: (message: string) => void,
+  addNewestConversation: (message: string, answer?: string) => void,
  removeLatestMessage: () => void,
+  addNewestAnswer: (answer: IAnswer) => void,
 ) => {
-  const loading = useOneNamespaceEffectsLoading('chatModel', [
-    'completeConversation',
-  ]);
  const { setConversation } = useSetConversation();
  const { conversationId } = useGetChatSearchParams();
  const { handleInputChange, value, setValue } = useHandleMessageInputChange();

  const fetchConversation = useFetchConversation();
-  const completeConversation = useCompleteConversation();

  const { handleClickConversation } = useClickConversationCard();
-  // const { send } = useConnectWithSseNext();
+  const { send, answer, done } = useSendMessageWithSse();

  const sendMessage = useCallback(
    async (message: string, id?: string) => {
-      const retcode = await completeConversation({
+      const res: Response = await send({
        conversation_id: id ?? conversationId,
        messages: [
          ...(conversation?.message ?? []).map((x: IMessage) => omit(x, 'id')),
@ -533,27 +570,33 @@ export const useSendMessage = (
        ],
      });

-      if (retcode === 0) {
+      if (res.status === 200) {
        if (id) {
+          console.info('111');
          // new conversation
          handleClickConversation(id);
        } else {
-          fetchConversation(conversationId);
+          console.info('222');
+          // fetchConversation(conversationId);
        }
      } else {
+        console.info('333');
+
        // cancel loading
        setValue(message);
+        console.info('removeLatestMessage111');
        removeLatestMessage();
      }
+      console.info('false');
    },
    [
      conversation?.message,
      conversationId,
-      fetchConversation,
+      // fetchConversation,
      handleClickConversation,
      removeLatestMessage,
      setValue,
-      completeConversation,
+      send,
    ],
  );

@ -572,19 +615,27 @@ export const useSendMessage = (
    [conversationId, setConversation, sendMessage],
  );

-  const handlePressEnter = () => {
-    if (!loading) {
+  useEffect(() => {
+    if (answer.answer) {
+      addNewestAnswer(answer);
+      console.info('true?');
+      console.info('send msg:', answer.answer);
+    }
+  }, [answer, addNewestAnswer]);
+
+  const handlePressEnter = useCallback(() => {
+    if (done) {
      setValue('');
-      addNewestConversation(value);
      handleSendMessage(value.trim());
    }
-  };
+    addNewestConversation(value);
+  }, [addNewestConversation, handleSendMessage, done, setValue, value]);

  return {
    handlePressEnter,
    handleInputChange,
    value,
-    loading,
+    loading: !done,
  };
 };

--- a/web/src/pages/chat/interface.ts
+++ b/web/src/pages/chat/interface.ts
@ -1,4 +1,4 @@
-import { IConversation, Message } from '@/interfaces/database/chat';
+import { IConversation, IReference, Message } from '@/interfaces/database/chat';
 import { FormInstance } from 'antd';

 export interface ISegmentedContentProps {
@ -24,6 +24,7 @@ export type IPromptConfigParameters = Omit<VariableTableDataType, 'variable'>;

 export interface IMessage extends Message {
  id: string;
+  reference?: IReference; // the latest news has reference
 }

 export interface IClientConversation extends IConversation {
--- a/web/src/pages/chat/markdown-content/index.less
+++ b/web/src/pages/chat/markdown-content/index.less
@ -23,3 +23,23 @@
 .referenceIcon {
  padding: 0 6px;
 }
+
+.cursor {
+  display: inline-block;
+  width: 1px;
+  height: 16px;
+  background-color: black;
+  animation: blink 0.6s infinite;
+  vertical-align: text-top;
+  @keyframes blink {
+    0% {
+      opacity: 1;
+    }
+    50% {
+      opacity: 0;
+    }
+    100% {
+      opacity: 1;
+    }
+  }
+}
--- a/web/src/pages/chat/markdown-content/index.tsx
+++ b/web/src/pages/chat/markdown-content/index.tsx
@ -16,6 +16,7 @@ import { visitParents } from 'unist-util-visit-parents';
 import styles from './index.less';

 const reg = /(#{2}\d+\${2})/g;
+const curReg = /(~{2}\d+\${2})/g;

 const getChunkIndex = (match: string) => Number(match.slice(2, -2));
 // TODO: The display of the table is inconsistent with the display previously placed in the MessageItem.
@ -61,7 +62,7 @@ const MarkdownContent = ({
    (chunkIndex: number) => {
      const chunks = reference?.chunks ?? [];
      const chunkItem = chunks[chunkIndex];
-      const document = reference?.doc_aggs.find(
+      const document = reference?.doc_aggs?.find(
        (x) => x?.doc_id === chunkItem?.doc_id,
      );
      const documentId = document?.doc_id;
@ -129,7 +130,7 @@ const MarkdownContent = ({

  const renderReference = useCallback(
    (text: string) => {
-      return reactStringReplace(text, reg, (match, i) => {
+      let replacedText = reactStringReplace(text, reg, (match, i) => {
        const chunkIndex = getChunkIndex(match);
        return (
          <Popover content={getPopoverContent(chunkIndex)}>
@ -137,6 +138,12 @@ const MarkdownContent = ({
          </Popover>
        );
      });
+
+      replacedText = reactStringReplace(replacedText, curReg, (match, i) => (
+        <span className={styles.cursor} key={i}></span>
+      ));
+
+      return replacedText;
    },
    [getPopoverContent],
  );
--- a/web/src/pages/chat/share/index.tsx
+++ b/web/src/pages/chat/share/index.tsx
@ -1,51 +1,11 @@
-import { useEffect } from 'react';
-import {
-  useCreateSharedConversationOnMount,
-  useSelectCurrentSharedConversation,
-  useSendSharedMessage,
-} from '../shared-hooks';
 import ChatContainer from './large';

 import styles from './index.less';

 const SharedChat = () => {
-  const { conversationId } = useCreateSharedConversationOnMount();
-  const {
-    currentConversation,
-    addNewestConversation,
-    removeLatestMessage,
-    ref,
-    loading,
-    setCurrentConversation,
-  } = useSelectCurrentSharedConversation(conversationId);
-
-  const {
-    handlePressEnter,
-    handleInputChange,
-    value,
-    loading: sendLoading,
-  } = useSendSharedMessage(
-    currentConversation,
-    addNewestConversation,
-    removeLatestMessage,
-    setCurrentConversation,
-  );
-
-  useEffect(() => {
-    console.info(location.href);
-  }, []);
-
  return (
    <div className={styles.chatWrapper}>
-      <ChatContainer
-        value={value}
-        handleInputChange={handleInputChange}
-        handlePressEnter={handlePressEnter}
-        loading={loading}
-        sendLoading={sendLoading}
-        conversation={currentConversation}
-        ref={ref}
-      ></ChatContainer>
+      <ChatContainer></ChatContainer>
    </div>
  );
 };
--- a/web/src/pages/chat/share/large.tsx
+++ b/web/src/pages/chat/share/large.tsx
@ -1,18 +1,50 @@
 import { ReactComponent as AssistantIcon } from '@/assets/svg/assistant.svg';
 import { MessageType } from '@/constants/chat';
 import { useTranslate } from '@/hooks/commonHooks';
-import { Message } from '@/interfaces/database/chat';
-import { Avatar, Button, Flex, Input, Skeleton, Spin } from 'antd';
+import { IReference, Message } from '@/interfaces/database/chat';
+import { Avatar, Button, Flex, Input, List, Spin } from 'antd';
 import classNames from 'classnames';
-import { useSelectConversationLoading } from '../hooks';

-import HightLightMarkdown from '@/components/highlight-markdown';
-import React, { ChangeEventHandler, forwardRef } from 'react';
-import { IClientConversation } from '../interface';
+import NewDocumentLink from '@/components/new-document-link';
+import SvgIcon from '@/components/svg-icon';
+import { useGetDocumentUrl } from '@/hooks/documentHooks';
+import { useSelectFileThumbnails } from '@/hooks/knowledgeHook';
+import { getExtension, isPdf } from '@/utils/documentUtils';
+import { forwardRef, useMemo } from 'react';
+import MarkdownContent from '../markdown-content';
+import {
+  useCreateSharedConversationOnMount,
+  useSelectCurrentSharedConversation,
+  useSendSharedMessage,
+} from '../shared-hooks';
+import { buildMessageItemReference } from '../utils';
 import styles from './index.less';

-const MessageItem = ({ item }: { item: Message }) => {
+const MessageItem = ({
+  item,
+  reference,
+  loading = false,
+}: {
+  item: Message;
+  reference: IReference;
+  loading?: boolean;
+}) => {
  const isAssistant = item.role === MessageType.Assistant;
+  const { t } = useTranslate('chat');
+  const fileThumbnails = useSelectFileThumbnails();
+  const getDocumentUrl = useGetDocumentUrl();
+
+  const referenceDocumentList = useMemo(() => {
+    return reference?.doc_aggs ?? [];
+  }, [reference?.doc_aggs]);
+
+  const content = useMemo(() => {
+    let text = item.content;
+    if (text === '') {
+      text = t('searching');
+    }
+    return loading ? text?.concat('~~2$$') : text;
+  }, [item.content, loading, t]);

  return (
    <div
@ -45,12 +77,43 @@ const MessageItem = ({ item }: { item: Message }) => {
          <Flex vertical gap={8} flex={1}>
            <b>{isAssistant ? '' : 'You'}</b>
            <div className={styles.messageText}>
-              {item.content !== '' ? (
-                <HightLightMarkdown>{item.content}</HightLightMarkdown>
-              ) : (
-                <Skeleton active className={styles.messageEmpty} />
-              )}
+              <MarkdownContent
+                reference={reference}
+                clickDocumentButton={() => {}}
+                content={content}
+              ></MarkdownContent>
            </div>
+            {isAssistant && referenceDocumentList.length > 0 && (
+              <List
+                bordered
+                dataSource={referenceDocumentList}
+                renderItem={(item) => {
+                  const fileThumbnail = fileThumbnails[item.doc_id];
+                  const fileExtension = getExtension(item.doc_name);
+                  return (
+                    <List.Item>
+                      <Flex gap={'small'} align="center">
+                        {fileThumbnail ? (
+                          <img src={fileThumbnail}></img>
+                        ) : (
+                          <SvgIcon
+                            name={`file-icon/${fileExtension}`}
+                            width={24}
+                          ></SvgIcon>
+                        )}
+
+                        <NewDocumentLink
+                          link={getDocumentUrl(item.doc_id)}
+                          preventDefault={!isPdf(item.doc_name)}
+                        >
+                          {item.doc_name}
+                        </NewDocumentLink>
+                      </Flex>
+                    </List.Item>
+                  );
+                }}
+              />
+            )}
          </Flex>
        </div>
      </section>
@ -58,28 +121,31 @@ const MessageItem = ({ item }: { item: Message }) => {
  );
 };

-interface IProps {
-  handlePressEnter(): void;
-  handleInputChange: ChangeEventHandler<HTMLInputElement>;
-  value: string;
-  loading: boolean;
-  sendLoading: boolean;
-  conversation: IClientConversation;
-  ref: React.LegacyRef<any>;
-}
+const ChatContainer = () => {
+  const { t } = useTranslate('chat');
+  const { conversationId } = useCreateSharedConversationOnMount();
+  const {
+    currentConversation: conversation,
+    addNewestConversation,
+    removeLatestMessage,
+    ref,
+    loading,
+    setCurrentConversation,
+    addNewestAnswer,
+  } = useSelectCurrentSharedConversation(conversationId);

-const ChatContainer = (
-  {
+  const {
    handlePressEnter,
    handleInputChange,
    value,
    loading: sendLoading,
+  } = useSendSharedMessage(
    conversation,
-  }: IProps,
-  ref: React.LegacyRef<any>,
-) => {
-  const loading = useSelectConversationLoading();
-  const { t } = useTranslate('chat');
+    addNewestConversation,
+    removeLatestMessage,
+    setCurrentConversation,
+    addNewestAnswer,
+  );

  return (
    <>
@ -87,9 +153,18 @@ const ChatContainer = (
        <Flex flex={1} vertical className={styles.messageContainer}>
          <div>
            <Spin spinning={loading}>
-              {conversation?.message?.map((message) => {
+              {conversation?.message?.map((message, i) => {
                return (
-                  <MessageItem key={message.id} item={message}></MessageItem>
+                  <MessageItem
+                    key={message.id}
+                    item={message}
+                    reference={buildMessageItemReference(conversation, message)}
+                    loading={
+                      message.role === MessageType.Assistant &&
+                      sendLoading &&
+                      conversation?.message.length - 1 === i
+                    }
+                  ></MessageItem>
                );
              })}
            </Spin>
--- a/web/src/pages/chat/shared-hooks.ts
+++ b/web/src/pages/chat/shared-hooks.ts
@ -1,10 +1,12 @@
 import { MessageType } from '@/constants/chat';
 import {
-  useCompleteSharedConversation,
  useCreateSharedConversation,
  useFetchSharedConversation,
 } from '@/hooks/chatHooks';
+import { useSendMessageWithSse } from '@/hooks/logicHooks';
 import { useOneNamespaceEffectsLoading } from '@/hooks/storeHooks';
+import { IAnswer } from '@/interfaces/database/chat';
+import api from '@/utils/api';
 import omit from 'lodash/omit';
 import {
  Dispatch,
@ -76,6 +78,27 @@ export const useSelectCurrentSharedConversation = (conversationId: string) => {
    });
  }, []);

+  const addNewestAnswer = useCallback((answer: IAnswer) => {
+    setCurrentConversation((pre) => {
+      const latestMessage = pre.message?.at(-1);
+
+      if (latestMessage) {
+        return {
+          ...pre,
+          message: [
+            ...pre.message.slice(0, -1),
+            {
+              ...latestMessage,
+              content: answer.answer,
+              reference: answer.reference,
+            } as IMessage,
+          ],
+        };
+      }
+      return pre;
+    });
+  }, []);
+
  const removeLatestMessage = useCallback(() => {
    setCurrentConversation((pre) => {
      const nextMessages = pre.message.slice(0, -2);
@ -106,6 +129,7 @@ export const useSelectCurrentSharedConversation = (conversationId: string) => {
    loading,
    ref,
    setCurrentConversation,
+    addNewestAnswer,
  };
 };

@ -114,20 +138,19 @@ export const useSendSharedMessage = (
  addNewestConversation: (message: string) => void,
  removeLatestMessage: () => void,
  setCurrentConversation: Dispatch<SetStateAction<IClientConversation>>,
+  addNewestAnswer: (answer: IAnswer) => void,
 ) => {
  const conversationId = conversation.id;
-  const loading = useOneNamespaceEffectsLoading('chatModel', [
-    'completeExternalConversation',
-  ]);
  const setConversation = useCreateSharedConversation();
  const { handleInputChange, value, setValue } = useHandleMessageInputChange();

-  const fetchConversation = useFetchSharedConversation();
-  const completeConversation = useCompleteSharedConversation();
+  const { send, answer, done } = useSendMessageWithSse(
+    api.completeExternalConversation,
+  );

  const sendMessage = useCallback(
    async (message: string, id?: string) => {
-      const retcode = await completeConversation({
+      const res: Response = await send({
        conversation_id: id ?? conversationId,
        quote: false,
        messages: [
@ -139,11 +162,11 @@ export const useSendSharedMessage = (
        ],
      });

-      if (retcode === 0) {
-        const data = await fetchConversation(conversationId);
-        if (data.retcode === 0) {
-          setCurrentConversation(data.data);
-        }
+      if (res?.status === 200) {
+        // const data = await fetchConversation(conversationId);
+        // if (data.retcode === 0) {
+        //   setCurrentConversation(data.data);
+        // }
      } else {
        // cancel loading
        setValue(message);
@ -153,11 +176,11 @@ export const useSendSharedMessage = (
    [
      conversationId,
      conversation?.message,
-      fetchConversation,
+      // fetchConversation,
      removeLatestMessage,
      setValue,
-      completeConversation,
-      setCurrentConversation,
+      send,
+      // setCurrentConversation,
    ],
  );

@ -176,18 +199,24 @@ export const useSendSharedMessage = (
    [conversationId, setConversation, sendMessage],
  );

-  const handlePressEnter = () => {
-    if (!loading) {
+  useEffect(() => {
+    if (answer.answer) {
+      addNewestAnswer(answer);
+    }
+  }, [answer, addNewestAnswer]);
+
+  const handlePressEnter = useCallback(() => {
+    if (done) {
      setValue('');
      addNewestConversation(value);
      handleSendMessage(value.trim());
    }
-  };
+  }, [addNewestConversation, done, handleSendMessage, setValue, value]);

  return {
    handlePressEnter,
    handleInputChange,
    value,
-    loading,
+    loading: !done,
  };
 };
--- a/web/src/pages/chat/utils.ts
+++ b/web/src/pages/chat/utils.ts
@ -1,5 +1,7 @@
+import { MessageType } from '@/constants/chat';
 import { IConversation, IReference } from '@/interfaces/database/chat';
 import { EmptyConversationId, variableEnabledFieldMap } from './constants';
+import { IClientConversation, IMessage } from './interface';

 export const excludeUnEnabledVariables = (values: any) => {
  const unEnabledFields: Array<keyof typeof variableEnabledFieldMap> =
@ -20,7 +22,7 @@ export const getDocumentIdsFromConversionReference = (data: IConversation) => {
  const documentIds = data.reference.reduce(
    (pre: Array<string>, cur: IReference) => {
      cur.doc_aggs
-        .map((x) => x.doc_id)
+        ?.map((x) => x.doc_id)
        .forEach((x) => {
          if (pre.every((y) => y !== x)) {
            pre.push(x);
@ -32,3 +34,20 @@ export const getDocumentIdsFromConversionReference = (data: IConversation) => {
  );
  return documentIds.join(',');
 };
+
+export const buildMessageItemReference = (
+  conversation: IClientConversation,
+  message: IMessage,
+) => {
+  const assistantMessages = conversation.message
+    ?.filter((x) => x.role === MessageType.Assistant)
+    .slice(1);
+  const referenceIndex = assistantMessages.findIndex(
+    (x) => x.id === message.id,
+  );
+  const reference = message?.reference
+    ? message?.reference
+    : conversation.reference[referenceIndex];
+
+  return reference;
+};