ragflow/web/src/pages/next-search/markdown-content/index.tsx

import Image from '@/components/image';
import SvgIcon from '@/components/svg-icon';
import { IReference, IReferenceChunk } from '@/interfaces/database/chat';
import { getExtension } from '@/utils/document-util';
import { InfoCircleOutlined } from '@ant-design/icons';
import DOMPurify from 'dompurify';
import React, { memo, useCallback, useEffect, useMemo } from 'react';
import Markdown from 'react-markdown';
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 { useTranslation } from 'react-i18next';

import 'katex/dist/katex.min.css'; // `rehype-katex` does not import the CSS for you

import ImageCarousel from '@/components/markdown-content/image-carousel';
import {
  groupConsecutiveReferences,
  shouldShowCarousel,
  type ReferenceGroup,
} from '@/components/markdown-content/reference-utils';
import {
  preprocessLaTeX,
  replaceTextByOldReg,
  replaceThinkToSection,
  showImage,
} from '@/utils/chat';

import { Button } from '@/components/ui/button';
import {
  Popover,
  PopoverContent,
  PopoverTrigger,
} from '@/components/ui/popover';
import { useFetchDocumentThumbnailsByIds } from '@/hooks/use-document-request';
import classNames from 'classnames';
import { omit } from 'lodash';
import { pipe } from 'lodash/fp';

// Defining Tailwind CSS class name constants
const styles = {
  referenceChunkImage: 'w-[10vw] object-contain',
  referenceInnerChunkImage: 'block object-contain max-w-full max-h-[6vh]',
  referenceImagePreview: 'max-w-[45vw] max-h-[45vh]',
  chunkContentText: 'max-h-[45vh] overflow-y-auto',
  documentLink: 'p-0',
  referenceIcon: 'px-[6px]',
  fileThumbnail: 'inline-block max-w-[40px]',
};

// TODO: The display of the table is inconsistent with the display previously placed in the MessageItem.
const MarkdownContent = ({
  reference,
  clickDocumentButton,
  content,
}: {
  content: string;
  loading: boolean;
  reference: IReference;
  clickDocumentButton?: (documentId: string, chunk: IReferenceChunk) => void;
}) => {
  const { t } = useTranslation();
  const { setDocumentIds, data: fileThumbnails } =
    useFetchDocumentThumbnailsByIds();
  const contentWithCursor = useMemo(() => {
    let text = DOMPurify.sanitize(content, {
      ADD_TAGS: ['think', 'section'],
      ADD_ATTR: ['class'],
    });
    // 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 = false,
      documentUrl?: string,
    ) =>
      () => {
        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 document = reference?.doc_aggs?.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 getPopoverContent = useCallback(
    (chunkIndex: number) => {
      const {
        fileThumbnail,
        fileExtension,
        imageId,
        chunkItem,
        documentId,
        document,
      } = getReferenceInfo(chunkIndex);

      return (
        <div key={chunkItem?.id} className="flex gap-2">
          {imageId && (
            <Popover>
              <PopoverTrigger>
                <Image
                  id={imageId}
                  className={styles.referenceChunkImage}
                ></Image>
              </PopoverTrigger>
              <PopoverContent>
                <Image
                  id={imageId}
                  className={styles.referenceImagePreview}
                ></Image>
              </PopoverContent>
            </Popover>
          )}
          <div className={'space-y-2 max-w-[40vw]'}>
            <div
              dangerouslySetInnerHTML={{
                __html: DOMPurify.sanitize(chunkItem?.content ?? ''),
              }}
              className={classNames(styles.chunkContentText)}
            ></div>
            {documentId && (
              <div className="flex gap-2">
                {fileThumbnail ? (
                  <img
                    src={fileThumbnail}
                    alt=""
                    className={styles.fileThumbnail}
                  />
                ) : (
                  <SvgIcon
                    name={`file-icon/${fileExtension}`}
                    width={24}
                  ></SvgIcon>
                )}
                <Button
                  variant="link"
                  className={classNames(styles.documentLink, 'text-wrap')}
                  onClick={handleDocumentButtonClick(
                    documentId,
                    chunkItem,
                    // fileExtension === 'pdf',
                    // documentUrl,
                  )}
                >
                  {document?.doc_name}
                </Button>
              </div>
            )}
          </div>
        </div>
      );
    },
    [getReferenceInfo, handleDocumentButtonClick],
  );

  const renderReference = useCallback(
    (text: string) => {
      const groups = groupConsecutiveReferences(text);
      const elements: React.ReactNode[] = [];
      let lastIndex = 0;

      groups.forEach((group: ReferenceGroup) => {
        // Add text before the group
        if (group[0].start > lastIndex) {
          elements.push(text.slice(lastIndex, group[0].start));
        }

        // Determine if this group should be a carousel
        if (shouldShowCarousel(group, reference)) {
          // Render carousel for consecutive image group
          elements.push(
            <ImageCarousel
              key={`carousel-${group[0].id}`}
              group={group}
              reference={reference}
              fileThumbnails={fileThumbnails}
              onImageClick={handleDocumentButtonClick}
            />,
          );
        } else {
          // Render individual references in the group
          group.forEach((ref) => {
            const chunkIndex = Number(ref.id);
            const { imageId, chunkItem, documentId } =
              getReferenceInfo(chunkIndex);
            const docType = chunkItem?.doc_type;

            if (showImage(docType)) {
              elements.push(
                <Image
                  key={ref.id}
                  id={imageId}
                  className={styles.referenceInnerChunkImage}
                  onClick={
                    documentId
                      ? handleDocumentButtonClick(documentId, chunkItem)
                      : () => {}
                  }
                />,
              );
            } else {
              elements.push(
                <Popover key={ref.id}>
                  <PopoverTrigger>
                    <InfoCircleOutlined className={styles.referenceIcon} />
                  </PopoverTrigger>
                  <PopoverContent className="!w-fit">
                    {getPopoverContent(chunkIndex)}
                  </PopoverContent>
                </Popover>,
              );
            }
            // Add the original reference text
            elements.push(ref.fullMatch);
          });
        }

        lastIndex = group[group.length - 1].end;
      });

      // Add any remaining text after the last group
      if (lastIndex < text.length) {
        elements.push(text.slice(lastIndex));
      }

      return elements;
    },
    [
      reference,
      fileThumbnails,
      handleDocumentButtonClick,
      getReferenceInfo,
      getPopoverContent,
    ],
  );

  return (
    <Markdown
      rehypePlugins={[rehypeWrapReference, rehypeKatex, rehypeRaw]}
      remarkPlugins={[remarkGfm, remarkMath]}
      className="[&>section.think]:pl-[10px] [&>section.think]:text-[#8b8b8b] [&>section.think]:border-l-2 [&>section.think]:border-l-[#d5d3d3] [&>section.think]:mb-[10px] [&>section.think]:text-xs [&>blockquote]:pl-[10px] [&>blockquote]:border-l-4 [&>blockquote]:border-l-[#ccc] text-sm"
      components={
        {
          'custom-typography': ({ children }: { children: string }) =>
            Array.isArray(renderReference(children))
              ? renderReference(children).map((element, index) => (
                  <React.Fragment key={index}>{element}</React.Fragment>
                ))
              : renderReference(children),
          code(props: any) {
            const { children, className, ...rest } = props;
            const restProps = omit(rest, 'node');
            const match = /language-(\w+)/.exec(className || '');
            return match ? (
              <SyntaxHighlighter
                {...restProps}
                PreTag="div"
                language={match[1]}
                wrapLongLines
              >
                {String(children).replace(/\n$/, '')}
              </SyntaxHighlighter>
            ) : (
              <code
                {...restProps}
                className={classNames(className, 'text-wrap')}
              >
                {children}
              </code>
            );
          },
        } as any
      }
    >
      {contentWithCursor}
    </Markdown>
  );
};

export default memo(MarkdownContent);