Feat: Display file references for agent dialogues #3221 (#8854)

### What problem does this PR solve?

Feat: Display file references for agent dialogues #3221
### Type of change


- [x] New Feature (non-breaking change which adds functionality)

web/src/components/next-markdown-content/index.less

@@ -0,0 +1,71 @@
+.markdownContentWrapper {
+  :global(section.think) {
+    padding-left: 10px;
+    color: #8b8b8b;
+    border-left: 2px solid #d5d3d3;
+    margin-bottom: 10px;
+    font-size: 12px;
+  }
+  :global(blockquote) {
+    padding-left: 10px;
+    border-left: 4px solid #ccc;
+  }
+}
+
+.referencePopoverWrapper {
+  max-width: 50vw;
+}
+
+.referenceChunkImage {
+  width: 10vw;
+  object-fit: contain;
+}
+
+.referenceInnerChunkImage {
+  display: block;
+  object-fit: contain;
+  max-width: 100%;
+  max-height: 6vh;
+}
+
+.referenceImagePreview {
+  max-width: 45vw;
+  max-height: 45vh;
+}
+.chunkContentText {
+  .chunkText;
+  max-height: 45vh;
+  overflow-y: auto;
+}
+.documentLink {
+  padding: 0;
+}
+
+.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;
+    }
+  }
+}
+
+.fileThumbnail {
+  display: inline-block;
+  max-width: 40px;
+}

web/src/components/next-markdown-content/index.tsx

@@ -0,0 +1,284 @@
+import Image from '@/components/image';
+import SvgIcon from '@/components/svg-icon';
+import { IReferenceChunk, IReferenceObject } from '@/interfaces/database/chat';
+import { getExtension } from '@/utils/document-util';
+import DOMPurify from 'dompurify';
+import { memo, useCallback, useEffect, useMemo } from 'react';
+import Markdown from 'react-markdown';
+import reactStringReplace from 'react-string-replace';
+import SyntaxHighlighter from 'react-syntax-highlighter';
+import rehypeKatex from 'rehype-katex';
+import rehypeRaw from 'rehype-raw';
+import remarkGfm from 'remark-gfm';
+import remarkMath from 'remark-math';
+import { visitParents } from 'unist-util-visit-parents';
+
+import { useFetchDocumentThumbnailsByIds } from '@/hooks/document-hooks';
+import { useTranslation } from 'react-i18next';
+
+import 'katex/dist/katex.min.css'; // `rehype-katex` does not import the CSS for you
+
+import {
+  preprocessLaTeX,
+  replaceThinkToSection,
+  showImage,
+} from '@/utils/chat';
+
+import { cn } from '@/lib/utils';
+import { currentReg, replaceTextByOldReg } from '@/pages/chat/utils';
+import classNames from 'classnames';
+import { pipe } from 'lodash/fp';
+import { CircleAlert } from 'lucide-react';
+import { Button } from '../ui/button';
+import {
+  HoverCard,
+  HoverCardContent,
+  HoverCardTrigger,
+} from '../ui/hover-card';
+import styles from './index.less';
+
+const getChunkIndex = (match: string) => Number(match);
+// TODO: The display of the table is inconsistent with the display previously placed in the MessageItem.
+function MarkdownContent({
+  reference,
+  clickDocumentButton,
+  content,
+}: {
+  content: string;
+  loading: boolean;
+  reference?: IReferenceObject;
+  clickDocumentButton?: (documentId: string, chunk: IReferenceChunk) => void;
+}) {
+  const { t } = useTranslation();
+  const { setDocumentIds, data: fileThumbnails } =
+    useFetchDocumentThumbnailsByIds();
+  const contentWithCursor = useMemo(() => {
+    // let text = DOMPurify.sanitize(content);
+    let text = content;
+    if (text === '') {
+      text = t('chat.searching');
+    }
+    const nextText = replaceTextByOldReg(text);
+    return pipe(replaceThinkToSection, preprocessLaTeX)(nextText);
+  }, [content, t]);
+
+  useEffect(() => {
+    const docAggs = reference?.doc_aggs;
+    setDocumentIds(Array.isArray(docAggs) ? docAggs.map((x) => x.doc_id) : []);
+  }, [reference, setDocumentIds]);
+
+  const handleDocumentButtonClick = useCallback(
+    (
+      documentId: string,
+      chunk: IReferenceChunk,
+      isPdf: boolean,
+      documentUrl?: string,
+    ) =>
+      () => {
+        if (!isPdf) {
+          if (!documentUrl) {
+            return;
+          }
+          window.open(documentUrl, '_blank');
+        } else {
+          clickDocumentButton?.(documentId, chunk);
+        }
+      },
+    [clickDocumentButton],
+  );
+
+  const rehypeWrapReference = () => {
+    return function wrapTextTransform(tree: any) {
+      visitParents(tree, 'text', (node, ancestors) => {
+        const latestAncestor = ancestors.at(-1);
+        if (
+          latestAncestor.tagName !== 'custom-typography' &&
+          latestAncestor.tagName !== 'code'
+        ) {
+          node.type = 'element';
+          node.tagName = 'custom-typography';
+          node.properties = {};
+          node.children = [{ type: 'text', value: node.value }];
+        }
+      });
+    };
+  };
+
+  const getReferenceInfo = useCallback(
+    (chunkIndex: number) => {
+      const chunks = reference?.chunks ?? {};
+      const chunkItem = chunks[chunkIndex];
+
+      const documentList = Object.values(reference?.doc_aggs ?? {});
+      const document = documentList.find(
+        (x) => x?.doc_id === chunkItem?.document_id,
+      );
+      const documentId = document?.doc_id;
+      const documentUrl = document?.url;
+      const fileThumbnail = documentId ? fileThumbnails[documentId] : '';
+      const fileExtension = documentId ? getExtension(document?.doc_name) : '';
+      const imageId = chunkItem?.image_id;
+
+      return {
+        documentUrl,
+        fileThumbnail,
+        fileExtension,
+        imageId,
+        chunkItem,
+        documentId,
+        document,
+      };
+    },
+    [fileThumbnails, reference],
+  );
+
+  const renderPopoverContent = useCallback(
+    (chunkIndex: number) => {
+      const {
+        documentUrl,
+        fileThumbnail,
+        fileExtension,
+        imageId,
+        chunkItem,
+        documentId,
+        document,
+      } = getReferenceInfo(chunkIndex);
+
+      return (
+        <div key={chunkItem?.id} className="flex gap-2">
+          {imageId && (
+            <HoverCard>
+              <HoverCardTrigger>
+                <Image
+                  id={imageId}
+                  className={styles.referenceChunkImage}
+                ></Image>
+              </HoverCardTrigger>
+              <HoverCardContent>
+                <Image
+                  id={imageId}
+                  className={cn(styles.referenceImagePreview)}
+                ></Image>
+              </HoverCardContent>
+            </HoverCard>
+          )}
+          <div className={'space-y-2 max-w-[40vw] w-full'}>
+            <div
+              dangerouslySetInnerHTML={{
+                __html: DOMPurify.sanitize(chunkItem?.content ?? ''),
+              }}
+              className={classNames(styles.chunkContentText, 'w-full')}
+            ></div>
+            {documentId && (
+              <div className="flex gap-1">
+                {fileThumbnail ? (
+                  <img
+                    src={fileThumbnail}
+                    alt=""
+                    className={styles.fileThumbnail}
+                  />
+                ) : (
+                  <SvgIcon
+                    name={`file-icon/${fileExtension}`}
+                    width={24}
+                  ></SvgIcon>
+                )}
+                <Button
+                  variant="link"
+                  onClick={handleDocumentButtonClick(
+                    documentId,
+                    chunkItem,
+                    fileExtension === 'pdf',
+                    documentUrl,
+                  )}
+                  className="text-ellipsis text-wrap"
+                >
+                  {document?.doc_name}
+                </Button>
+              </div>
+            )}
+          </div>
+        </div>
+      );
+    },
+    [getReferenceInfo, handleDocumentButtonClick],
+  );
+
+  const renderReference = useCallback(
+    (text: string) => {
+      let replacedText = reactStringReplace(text, currentReg, (match, i) => {
+        const chunkIndex = getChunkIndex(match);
+
+        const { documentUrl, fileExtension, imageId, chunkItem, documentId } =
+          getReferenceInfo(chunkIndex);
+
+        const docType = chunkItem?.doc_type;
+
+        return showImage(docType) ? (
+          <Image
+            id={imageId}
+            className={styles.referenceInnerChunkImage}
+            onClick={
+              documentId
+                ? handleDocumentButtonClick(
+                    documentId,
+                    chunkItem,
+                    fileExtension === 'pdf',
+                    documentUrl,
+                  )
+                : () => {}
+            }
+          ></Image>
+        ) : (
+          <HoverCard key={i}>
+            <HoverCardTrigger>
+              <CircleAlert className="size-4 inline-block" />
+            </HoverCardTrigger>
+            <HoverCardContent>
+              {renderPopoverContent(chunkIndex)}
+            </HoverCardContent>
+          </HoverCard>
+        );
+      });
+
+      return replacedText;
+    },
+    [renderPopoverContent, getReferenceInfo, handleDocumentButtonClick],
+  );
+
+  return (
+    <Markdown
+      rehypePlugins={[rehypeWrapReference, rehypeKatex, rehypeRaw]}
+      remarkPlugins={[remarkGfm, remarkMath]}
+      className={styles.markdownContentWrapper}
+      components={
+        {
+          'custom-typography': ({ children }: { children: string }) =>
+            renderReference(children),
+          code(props: any) {
+            const { children, className, node, ...rest } = props;
+            const match = /language-(\w+)/.exec(className || '');
+            return match ? (
+              <SyntaxHighlighter
+                {...rest}
+                PreTag="div"
+                language={match[1]}
+                wrapLongLines
+              >
+                {String(children).replace(/\n$/, '')}
+              </SyntaxHighlighter>
+            ) : (
+              <code {...rest} className={classNames(className, 'text-wrap')}>
+                {children}
+              </code>
+            );
+          },
+        } as any
+      }
+    >
+      {contentWithCursor}
+    </Markdown>
+  );
+}
+
+export default memo(MarkdownContent);

web/src/components/next-message-item/index.less

@@ -30,7 +30,6 @@
   .messageTextDark {
     .chunkText();
     .messageTextBase();
-    background-color: #1668dc;
     word-break: break-word;
     :global(section.think) {
       color: rgb(166, 166, 166);
@@ -41,7 +40,6 @@
   .messageUserText {
     .chunkText();
     .messageTextBase();
-    background-color: rgba(255, 255, 255, 0.3);
     word-break: break-word;
     text-align: justify;
   }

web/src/components/next-message-item/index.tsx

@@ -1,7 +1,7 @@
 import { ReactComponent as AssistantIcon } from '@/assets/svg/assistant.svg';
 import { MessageType } from '@/constants/chat';
 import { useSetModalState } from '@/hooks/common-hooks';
-import { IReference, IReferenceChunk } from '@/interfaces/database/chat';
+import { IReferenceChunk, IReferenceObject } from '@/interfaces/database/chat';
 import classNames from 'classnames';
 import {
   PropsWithChildren,
@@ -17,16 +17,19 @@ import {
   useFetchDocumentThumbnailsByIds,
 } from '@/hooks/document-hooks';
 import { IRegenerateMessage, IRemoveMessageById } from '@/hooks/logic-hooks';
+import { cn } from '@/lib/utils';
 import { IMessage } from '@/pages/chat/interface';
-import MarkdownContent from '@/pages/chat/markdown-content';
 import { getExtension, isImage } from '@/utils/document-util';
 import { Avatar, Button, Flex, List, Space, Typography } from 'antd';
+import { isEmpty } from 'lodash';
 import FileIcon from '../file-icon';
 import IndentedTreeModal from '../indented-tree/modal';
 import NewDocumentLink from '../new-document-link';
+import MarkdownContent from '../next-markdown-content';
 import { useTheme } from '../theme-provider';
 import { AssistantGroupButton, UserGroupButton } from './group-button';
 import styles from './index.less';
+import { ReferenceDocumentList } from './reference-document-list';
 
 const { Text } = Typography;
 
@@ -35,7 +38,7 @@ interface IProps
     IRegenerateMessage,
     PropsWithChildren {
   item: IMessage;
-  reference: IReference;
+  reference?: IReferenceObject;
   loading?: boolean;
   sendLoading?: boolean;
   visibleAvatar?: boolean;
@@ -48,7 +51,7 @@ interface IProps
   showLoudspeaker?: boolean;
 }
 
-const MessageItem = ({
+function MessageItem({
   item,
   reference,
   loading = false,
@@ -56,14 +59,13 @@ const MessageItem = ({
   avatarDialog,
   sendLoading = false,
   clickDocumentButton,
-  index,
   removeMessageById,
   regenerateMessage,
   showLikeButton = true,
   showLoudspeaker = true,
   visibleAvatar = true,
   children,
-}: IProps) => {
+}: IProps) {
   const { theme } = useTheme();
   const isAssistant = item.role === MessageType.Assistant;
   const isUser = item.role === MessageType.User;
@@ -73,8 +75,10 @@ const MessageItem = ({
   const { visible, hideModal, showModal } = useSetModalState();
   const [clickedDocumentId, setClickedDocumentId] = useState('');
 
-  const referenceDocumentList = useMemo(() => {
-    return reference?.doc_aggs ?? [];
+  const referenceDocuments = useMemo(() => {
+    const docs = reference?.doc_aggs ?? {};
+
+    return Object.values(docs);
   }, [reference?.doc_aggs]);
 
   const handleUserDocumentClick = useCallback(
@@ -153,16 +157,18 @@ const MessageItem = ({
               {/* <b>{isAssistant ? '' : nickname}</b> */}
             </Space>
             <div
-              className={
-                isAssistant
-                  ? theme === 'dark'
-                    ? styles.messageTextDark
-                    : styles.messageText
-                  : styles.messageUserText
-              }
+              className={cn({
+                [theme === 'dark'
+                  ? styles.messageTextDark
+                  : styles.messageText]: isAssistant,
+                [styles.messageUserText]: !isAssistant,
+                'bg-background-card': !isAssistant,
+              })}
             >
               {item.data ? (
                 children
+              ) : sendLoading && isEmpty(item.content) ? (
+                'searching...'
               ) : (
                 <MarkdownContent
                   loading={loading}
@@ -172,32 +178,10 @@ const MessageItem = ({
                 ></MarkdownContent>
               )}
             </div>
-            {isAssistant && referenceDocumentList.length > 0 && (
-              <List
-                bordered
-                dataSource={referenceDocumentList}
-                renderItem={(item) => {
-                  return (
-                    <List.Item>
-                      <Flex gap={'small'} align="center">
-                        <FileIcon
-                          id={item.doc_id}
-                          name={item.doc_name}
-                        ></FileIcon>
-
-                        <NewDocumentLink
-                          documentId={item.doc_id}
-                          documentName={item.doc_name}
-                          prefix="document"
-                          link={item.url}
-                        >
-                          {item.doc_name}
-                        </NewDocumentLink>
-                      </Flex>
-                    </List.Item>
-                  );
-                }}
-              />
+            {isAssistant && referenceDocuments.length > 0 && (
+              <ReferenceDocumentList
+                list={referenceDocuments}
+              ></ReferenceDocumentList>
             )}
             {isUser && documentList.length > 0 && (
               <List
@@ -252,6 +236,6 @@ const MessageItem = ({
       )}
     </div>
   );
-};
+}
 
 export default memo(MessageItem);

web/src/components/next-message-item/reference-document-list.tsx

@@ -0,0 +1,27 @@
+import { Card, CardContent } from '@/components/ui/card';
+import { Docagg } from '@/interfaces/database/chat';
+import FileIcon from '../file-icon';
+import NewDocumentLink from '../new-document-link';
+
+export function ReferenceDocumentList({ list }: { list: Docagg[] }) {
+  return (
+    <section className="flex gap-3 flex-wrap">
+      {list.map((item) => (
+        <Card key={item.doc_id}>
+          <CardContent className="p-2">
+            <FileIcon id={item.doc_id} name={item.doc_name}></FileIcon>
+            <NewDocumentLink
+              documentId={item.doc_id}
+              documentName={item.doc_name}
+              prefix="document"
+              link={item.url}
+              className="text-text-sub-title-invert"
+            >
+              {item.doc_name}
+            </NewDocumentLink>
+          </CardContent>
+        </Card>
+      ))}
+    </section>
+  );
+}

web/src/hooks/use-send-message.ts

@@ -1,4 +1,5 @@
 import { Authorization } from '@/constants/authorization';
+import { IReferenceObject } from '@/interfaces/database/chat';
 import { BeginQuery } from '@/pages/agent/interface';
 import api from '@/utils/api';
 import { getAuthorization } from '@/utils/authorization-util';
@@ -43,6 +44,10 @@ export interface IMessageData {
   content: string;
 }
 
+export interface IMessageEndData {
+  reference: IReferenceObject;
+}
+
 export interface ILogData extends INodeData {
   logs: {
     name: string;
@@ -58,11 +63,13 @@ export type INodeEvent = IAnswerEvent<INodeData>;
 
 export type IMessageEvent = IAnswerEvent<IMessageData>;
 
+export type IMessageEndEvent = IAnswerEvent<IMessageEndData>;
+
 export type IInputEvent = IAnswerEvent<IInputData>;
 
 export type ILogEvent = IAnswerEvent<ILogData>;
 
-export type IChatEvent = INodeEvent | IMessageEvent;
+export type IChatEvent = INodeEvent | IMessageEvent | IMessageEndEvent;
 
 export type IEventList = Array<IChatEvent>;
 

web/src/interfaces/database/chat.ts

@@ -96,6 +96,11 @@ export interface IReference {
   total: number;
 }
 
+export interface IReferenceObject {
+  chunks: Record<string, IReferenceChunk>;
+  doc_aggs: Record<string, Docagg>;
+}
+
 export interface IAnswer {
   answer: string;
   reference?: IReference;

web/src/pages/agent/chat/box.tsx

@@ -18,7 +18,6 @@ import { useParams } from 'umi';
 import DebugContent from '../debug-content';
 import { BeginQuery } from '../interface';
 import { buildBeginQueryWithObject } from '../utils';
-import { buildAgentMessageItemReference } from '../utils/chat';
 
 const AgentChatBox = () => {
   const {
@@ -29,9 +28,9 @@ const AgentChatBox = () => {
     loading,
     ref,
     derivedMessages,
-    reference,
     stopOutputMessage,
     sendFormMessage,
+    findReferenceByMessageId,
   } = useSendNextMessage();
 
   const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
@@ -71,7 +70,7 @@ const AgentChatBox = () => {
 
   return (
     <>
-      <section className="flex flex-1 flex-col pl-5 h-[90vh]">
+      <section className="flex flex-1 flex-col px-5 h-[90vh]">
         <div className="flex-1 overflow-auto">
           <div>
             <Spin spinning={loading}>
@@ -88,10 +87,7 @@ const AgentChatBox = () => {
                     avatar={userInfo.avatar}
                     avatarDialog={canvasInfo.avatar}
                     item={message}
-                    reference={buildAgentMessageItemReference(
-                      { message: derivedMessages, reference },
-                      message,
-                    )}
+                    reference={findReferenceByMessageId(message.id)}
                     clickDocumentButton={clickDocumentButton}
                     index={i}
                     showLikeButton={false}

web/src/pages/agent/chat/chat-sheet.tsx

@@ -1,24 +1,18 @@
-import {
-  Sheet,
-  SheetContent,
-  SheetHeader,
-  SheetTitle,
-} from '@/components/ui/sheet';
+import { Sheet, SheetContent } from '@/components/ui/sheet';
 import { IModalProps } from '@/interfaces/common';
 import { cn } from '@/lib/utils';
+import { useTranslation } from 'react-i18next';
 import AgentChatBox from './box';
 
 export function ChatSheet({ hideModal }: IModalProps<any>) {
+  const { t } = useTranslation();
   return (
     <Sheet open modal={false} onOpenChange={hideModal}>
-      <SheetTitle className="hidden"></SheetTitle>
       <SheetContent
         className={cn('top-20 p-0')}
         onInteractOutside={(e) => e.preventDefault()}
       >
-        <SheetHeader>
-          <SheetTitle>Are you absolutely sure?</SheetTitle>
-        </SheetHeader>
+        <div className="pl-5 pt-2">{t('chat.chat')}</div>
         <AgentChatBox></AgentChatBox>
       </SheetContent>
     </Sheet>

web/src/pages/agent/chat/hooks.ts

@@ -8,6 +8,8 @@ import { useFetchAgent } from '@/hooks/use-agent-request';
 import {
   IEventList,
   IInputEvent,
+  IMessageEndData,
+  IMessageEndEvent,
   IMessageEvent,
   MessageEventType,
   useSendMessageBySSE,
@@ -17,7 +19,7 @@ import i18n from '@/locales/config';
 import api from '@/utils/api';
 import { get } from 'lodash';
 import trim from 'lodash/trim';
-import { useCallback, useContext, useEffect, useMemo } from 'react';
+import { useCallback, useContext, useEffect, useMemo, useState } from 'react';
 import { useParams } from 'umi';
 import { v4 as uuid } from 'uuid';
 import { BeginId } from '../constant';
@@ -114,6 +116,9 @@ export const useSendNextMessage = () => {
   const { refetch } = useFetchAgent();
   const { addEventList } = useContext(AgentChatLogContext);
   const getBeginNodeDataQuery = useGetBeginNodeDataQuery();
+  const [messageEndEventList, setMessageEndEventList] = useState<
+    IMessageEndEvent[]
+  >([]);
 
   const { send, answerList, done, stopOutputMessage } = useSendMessageBySSE(
     api.runCanvas,
@@ -126,13 +131,16 @@ export const useSendNextMessage = () => {
       const params: Record<string, unknown> = {
         id: agentId,
       };
+
       params.running_hint_text = i18n.t('flow.runningHintText', {
         defaultValue: 'is running...🕞',
       });
       if (message.content) {
+        const query = getBeginNodeDataQuery();
+
         params.query = message.content;
         // params.message_id = message.id;
-        params.inputs = {}; // begin operator inputs
+        params.inputs = transferInputsArrayToObject(query); // begin operator inputs
       }
       const res = await send(params);
 
@@ -146,7 +154,14 @@ export const useSendNextMessage = () => {
         refetch(); // pull the message list after sending the message successfully
       }
     },
-    [agentId, send, setValue, removeLatestMessage, refetch],
+    [
+      agentId,
+      send,
+      getBeginNodeDataQuery,
+      setValue,
+      removeLatestMessage,
+      refetch,
+    ],
   );
 
   const handleSendMessage = useCallback(
@@ -156,6 +171,23 @@ export const useSendNextMessage = () => {
     [sendMessage],
   );
 
+  useEffect(() => {
+    const messageEndEvent = answerList.find(
+      (x) => x.event === MessageEventType.MessageEnd,
+    );
+    if (messageEndEvent) {
+      setMessageEndEventList((list) => {
+        const nextList = [...list];
+        if (
+          nextList.every((x) => x.message_id !== messageEndEvent.message_id)
+        ) {
+          nextList.push(messageEndEvent as IMessageEndEvent);
+        }
+        return nextList;
+      });
+    }
+  }, [addEventList.length, answerList]);
+
   useEffect(() => {
     const { content, id } = findMessageFromList(answerList);
     const inputAnswer = findInputFromList(answerList);
@@ -195,11 +227,20 @@ export const useSendNextMessage = () => {
     [addNewestOneQuestion, send],
   );
 
+  const findReferenceByMessageId = useCallback(
+    (messageId: string) => {
+      const event = messageEndEventList.find(
+        (item) => item.message_id === messageId,
+      );
+      if (event) {
+        return (event?.data as IMessageEndData)?.reference;
+      }
+    },
+    [messageEndEventList],
+  );
+
   useEffect(() => {
-    const query = getBeginNodeDataQuery();
-    if (query.length > 0) {
-      send({ id: agentId, inputs: transferInputsArrayToObject(query) });
-    } else if (prologue) {
+    if (prologue) {
       addNewestOneAnswer({
         answer: prologue,
       });
@@ -230,5 +271,6 @@ export const useSendNextMessage = () => {
     stopOutputMessage,
     send,
     sendFormMessage,
+    findReferenceByMessageId,
   };
 };