jk-rag-chat-client/app/components/markdown.tsx

import ReactMarkdown from "react-markdown";
import { Image } from "antd";
import RemarkMath from "remark-math";
import RemarkBreaks from "remark-breaks";
import RehypeKatex from "rehype-katex";

// import rehypeMathjaxChtml from 'rehype-mathjax/chtml';

import RemarkGfm from "remark-gfm";
import RehypeHighlight from "rehype-highlight";
import { useRef, useState, RefObject, useEffect, useMemo } from "react";
import { copyToClipboard, useWindowSize } from "../utils";
import mermaid from "mermaid";
import LoadingIcon from "../icons/three-dots.svg";
import React from "react";
import { useDebouncedCallback } from "use-debounce";
import { showImageModal, FullScreen } from "./ui-lib";
import { ArtifactsShareButton, HTMLPreview } from "./artifacts";
import { Plugin } from "../constant";
import { useChatStore } from "../store";
import "katex/dist/katex.min.css";

export function Mermaid(props: { code: string }) {
  const ref = useRef<HTMLDivElement>(null);
  const [hasError, setHasError] = useState(false);

  useEffect(() => {
    if (props.code && ref.current) {
      mermaid
        .run({
          nodes: [ref.current],
          suppressErrors: true,
        })
        .catch((e) => {
          setHasError(true);
          console.error("[Mermaid] ", e.message);
        });
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [props.code]);

  function viewSvgInNewWindow() {
    const svg = ref.current?.querySelector("svg");
    if (!svg) return;
    const text = new XMLSerializer().serializeToString(svg);
    const blob = new Blob([text], { type: "image/svg+xml" });
    showImageModal(URL.createObjectURL(blob));
  }

  if (hasError) {
    return null;
  }

  return (
    <div
      className="no-dark mermaid"
      style={{
        cursor: "pointer",
        overflow: "auto",
      }}
      ref={ref}
      onClick={() => viewSvgInNewWindow()}
    >
      {props.code}
    </div>
  );
}

export function PreCode(props: { children: any }) {
  const ref = useRef<HTMLPreElement>(null);
  const refText = ref.current?.innerText;
  const [mermaidCode, setMermaidCode] = useState("");
  const [htmlCode, setHtmlCode] = useState("");
  const { height } = useWindowSize();
  const chatStore = useChatStore();
  const session = chatStore.currentSession();
  const plugins = session.mask?.plugin;

  const renderArtifacts = useDebouncedCallback(() => {
    if (!ref.current) return;
    const mermaidDom = ref.current.querySelector("code.language-mermaid");
    if (mermaidDom) {
      setMermaidCode((mermaidDom as HTMLElement).innerText);
    }
    const htmlDom = ref.current.querySelector("code.language-html");
    if (htmlDom) {
      setHtmlCode((htmlDom as HTMLElement).innerText);
    } else if (refText?.startsWith("<!DOCTYPE")) {
      setHtmlCode(refText);
    }
  }, 600);

  useEffect(() => {
    setTimeout(renderArtifacts, 1);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [refText]);

  const enableArtifacts = useMemo(
    () => plugins?.includes(Plugin.Artifacts),
    [plugins],
  );

  //Wrap the paragraph for plain-text
  useEffect(() => {
    if (ref.current) {
      const codeElements = ref.current.querySelectorAll(
        "code",
      ) as NodeListOf<HTMLElement>;
      const wrapLanguages = [
        "",
        "think",
        "md",
        "markdown",
        "text",
        "txt",
        "plaintext",
        "tex",
        "latex",
      ];
      codeElements.forEach((codeElement) => {
        let languageClass = codeElement.className.match(/language-(\w+)/);
        let name = languageClass ? languageClass[1] : "";
        if (wrapLanguages.includes(name)) {
          codeElement.style.whiteSpace = "pre-wrap";
        }
      });
    }
  }, []);

  return (
    <>
      <pre ref={ref}>
        <span
          className="copy-code-button"
          onClick={() => {
            if (ref.current) {
              const code = ref.current.innerText;
              copyToClipboard(code);
            }
          }}
        ></span>
        {props.children}
      </pre>
      {mermaidCode.length > 0 && (
        <Mermaid code={mermaidCode} key={mermaidCode} />
      )}
      {htmlCode.length > 0 && enableArtifacts && (
        <FullScreen className="no-dark html" right={70}>
          <ArtifactsShareButton
            style={{ position: "absolute", right: 20, top: 10 }}
            getCode={() => htmlCode}
          />
          <HTMLPreview
            code={htmlCode}
            autoHeight={!document.fullscreenElement}
            height={!document.fullscreenElement ? 600 : height}
          />
        </FullScreen>
      )}
    </>
  );
}

function escapeDollarNumber(text: string) {
  // 原有的 escapeDollarNumber 逻辑被移除，因为它会破坏以数字开头的公式
  // 如果需要防止 $100 被识别为公式，应建议用户正确转义 \$100，或优化 preprocessLaTeX 的识别逻辑
  return text;
}

function preprocessLaTeX(content: string) {
  const pattern =
    /(```[\s\S]*?```|`.*?`)|\\\[([\s\S]*?[^\\])\\\]|\\\((.*?)\\\)|(\$\$[\s\S]*?\$\$)|(\$(?!\s)[\s\S]*?\S\$)/g;
  
  return content.replace(
    pattern,
    (match, codeBlock, squareBracket, roundBracket, doubleDollar, singleDollar) => {
      if (codeBlock) {
        return codeBlock;
      }

      let innerContent = "";
      let isBlock = false;

      if (squareBracket) {
        innerContent = squareBracket;
        isBlock = true;
      } else if (roundBracket) {
        innerContent = roundBracket;
        isBlock = false;
      } else if (doubleDollar) {
        innerContent = doubleDollar.slice(2, -2);
        isBlock = true;
      } else if (singleDollar) {
        innerContent = singleDollar.slice(1, -1);
        isBlock = false;
      } else {
        return match;
      }

      // 修复1: \200,000 -> 200,000 (移除数字前的反斜杠)
      innerContent = innerContent.replace(/\\(\d)/g, "$1");
      // 修复2: 400;kg -> 400\;kg (分号间隔)
      innerContent = innerContent.replace(/(\d+)\s*;\s*/g, "$1\\;");

      return isBlock ? `$$${innerContent}$$` : `$${innerContent}$`;
    },
  );
}

function _MarkDownContent(props: { content: string }) {
  const escapedContent = useMemo(() => {
    return preprocessLaTeX(props.content);
  }, [props.content]);

  return (
    <ReactMarkdown
      remarkPlugins={[RemarkMath, RemarkGfm, RemarkBreaks]}
      rehypePlugins={[
        [RehypeKatex,
          {
            strict: 'ignore',
            throwOnError: false, // 关键：即使有错误也不抛出异常，仅渲染错误内容
            errorColor: '#ff0000', // 错误内容标红（StackEdit同款）
            trust: true, // 允许更多语法，提升兼容性
          }],
        [
          RehypeHighlight,
          {
            detect: false,
            ignoreMissing: true,
          },
        ],
      ]}
      // 控制不同标签的显示样式
      components={{
        pre: PreCode,
        head: () => (
          <link
            rel="stylesheet"
            href="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.css"
            integrity="sha384-n8MVd4RsNIU0tAv4ct0nTaAbDJwPJzDEaqSD1odI+WdtXRGWt2kTvGFasHpSy3SV"
            crossOrigin="anonymous"
          />
        ),
        code: ({ className, children }) => {
          // console.log('className, children-----3',className, children)
          if (className && className.includes('language-think')) {
            return (
              <code style={{ whiteSpace: 'pre-wrap', background: '#f3f4f6', color: '#525252' }}>
                {children}
              </code>
            );
          } else {
            return children;
          }
        },
        div: (pProps) => {
          // console.log('pProps-----2',pProps)
          return <div {...pProps} dir="auto" />
        },
        a: (aProps) => {
          // console.log('aProps-----4',aProps)
          const href = aProps.href || "";
          const isInternal = /^\/#/i.test(href);
          const target = isInternal ? "_self" : aProps.target ?? "_blank";
          return <a {...aProps} target={target} />;
        },
        img: ({ src, alt }) => (
          // console.log('src, alt-----5',src, alt),
          <span style={{ width: '100%', height: 'auto', cursor: 'pointer', display: 'inline-block' }}>
            <Image
              width='80%'
              src={src}
              alt={alt}
              preview={{
                mask: null
              }}
            />
          </span>
        ),
      }}
    >
      {escapedContent}
    </ReactMarkdown >
  );
}

export const MarkdownContent = React.memo(_MarkDownContent);

export function Markdown(
  props: {
    content: string;
    loading?: boolean;
    fontSize?: number;
    fontFamily?: string;
    parentRef?: RefObject<HTMLDivElement>;
    defaultShow?: boolean;
  } & React.DOMAttributes<HTMLDivElement>,
) {
  const mdRef = useRef<HTMLDivElement>(null);

  return (
    <div
      className="markdown-body"
      style={{
        fontSize: `${props.fontSize ?? 14}px`,
        fontFamily: props.fontFamily || "inherit",
      }}
      ref={mdRef}
      onContextMenu={props.onContextMenu}
      onDoubleClickCapture={props.onDoubleClickCapture}
      dir="auto"
    >
      {props.loading ? (
        <LoadingIcon />
      ) : (
        <MarkdownContent content={props.content} />
      )}
    </div>
  );
}