KBaba7/llama.cpp

1import React, { useMemo, useState } from 'react';2import Markdown, { ExtraProps } from 'react-markdown';3import remarkGfm from 'remark-gfm';4import rehypeHightlight from 'rehype-highlight';5import rehypeKatex from 'rehype-katex';6import remarkMath from 'remark-math';7import remarkBreaks from 'remark-breaks';8import 'katex/dist/katex.min.css';9import { classNames, copyStr } from '../utils/misc';10import { ElementContent, Root } from 'hast';11import { visit } from 'unist-util-visit';12 13export default function MarkdownDisplay({ content }: { content: string }) {14  const preprocessedContent = useMemo(15    () => preprocessLaTeX(content),16    [content]17  );18  return (19    <Markdown20      remarkPlugins={[remarkGfm, remarkMath, remarkBreaks]}21      rehypePlugins={[rehypeHightlight, rehypeKatex, rehypeCustomCopyButton]}22      components={{23        button: (props) => (24          <CopyCodeButton {...props} origContent={preprocessedContent} />25        ),26      }}27    >28      {preprocessedContent}29    </Markdown>30  );31}32 33const CopyCodeButton: React.ElementType<34  React.ClassAttributes<HTMLButtonElement> &35    React.HTMLAttributes<HTMLButtonElement> &36    ExtraProps & { origContent: string }37> = ({ node, origContent }) => {38  const startOffset = node?.position?.start.offset ?? 0;39  const endOffset = node?.position?.end.offset ?? 0;40 41  const copiedContent = useMemo(42    () =>43      origContent44        .substring(startOffset, endOffset)45        .replace(/^```[^\n]+\n/g, '')46        .replace(/```$/g, ''),47    [origContent, startOffset, endOffset]48  );49 50  return (51    <div52      className={classNames({53        'text-right sticky top-4 mb-2 mr-2 h-0': true,54        'display-none': !node?.position,55      })}56    >57      <CopyButton className="badge btn-mini" content={copiedContent} />58    </div>59  );60};61 62export const CopyButton = ({63  content,64  className,65}: {66  content: string;67  className?: string;68}) => {69  const [copied, setCopied] = useState(false);70  return (71    <button72      className={className}73      onClick={() => {74        copyStr(content);75        setCopied(true);76      }}77      onMouseLeave={() => setCopied(false)}78    >79      {copied ? 'Copied!' : '📋 Copy'}80    </button>81  );82};83 84/**85 * This injects the "button" element before each "pre" element.86 * The actual button will be replaced with a react component in the MarkdownDisplay.87 * We don't replace "pre" node directly because it will cause the node to re-render, which causes this bug: https://github.com/ggerganov/llama.cpp/issues/960888 */89function rehypeCustomCopyButton() {90  return function (tree: Root) {91    visit(tree, 'element', function (node) {92      if (node.tagName === 'pre' && !node.properties.visited) {93        const preNode = { ...node };94        // replace current node95        preNode.properties.visited = 'true';96        node.tagName = 'div';97        node.properties = {98          className: 'relative my-4',99        };100        // add node for button101        const btnNode: ElementContent = {102          type: 'element',103          tagName: 'button',104          properties: {},105          children: [],106          position: node.position,107        };108        node.children = [btnNode, preNode];109      }110    });111  };112}113 114/**115 * The part below is copied and adapted from:116 * https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts117 * (MIT License)118 */119 120// Regex to check if the processed content contains any potential LaTeX patterns121const containsLatexRegex =122  /\\\(.*?\\\)|\\\[.*?\\\]|\$.*?\$|\\begin\{equation\}.*?\\end\{equation\}/;123 124// Regex for inline and block LaTeX expressions125const inlineLatex = new RegExp(/\\\((.+?)\\\)/, 'g');126const blockLatex = new RegExp(/\\\[(.*?[^\\])\\\]/, 'gs');127 128// Function to restore code blocks129const restoreCodeBlocks = (content: string, codeBlocks: string[]) => {130  return content.replace(131    /<<CODE_BLOCK_(\d+)>>/g,132    (_, index) => codeBlocks[index]133  );134};135 136// Regex to identify code blocks and inline code137const codeBlockRegex = /(```[\s\S]*?```|`.*?`)/g;138 139export const processLaTeX = (_content: string) => {140  let content = _content;141  // Temporarily replace code blocks and inline code with placeholders142  const codeBlocks: string[] = [];143  let index = 0;144  content = content.replace(codeBlockRegex, (match) => {145    codeBlocks[index] = match;146    return `<<CODE_BLOCK_${index++}>>`;147  });148 149  // Escape dollar signs followed by a digit or space and digit150  let processedContent = content.replace(/(\$)(?=\s?\d)/g, '\\$');151 152  // If no LaTeX patterns are found, restore code blocks and return the processed content153  if (!containsLatexRegex.test(processedContent)) {154    return restoreCodeBlocks(processedContent, codeBlocks);155  }156 157  // Convert LaTeX expressions to a markdown compatible format158  processedContent = processedContent159    .replace(inlineLatex, (_: string, equation: string) => `$${equation}$`) // Convert inline LaTeX160    .replace(blockLatex, (_: string, equation: string) => `$$${equation}$$`); // Convert block LaTeX161 162  // Restore code blocks163  return restoreCodeBlocks(processedContent, codeBlocks);164};165 166/**167 * Preprocesses LaTeX content by replacing delimiters and escaping certain characters.168 *169 * @param content The input string containing LaTeX expressions.170 * @returns The processed string with replaced delimiters and escaped characters.171 */172export function preprocessLaTeX(content: string): string {173  // Step 1: Protect code blocks174  const codeBlocks: string[] = [];175  content = content.replace(/(```[\s\S]*?```|`[^`\n]+`)/g, (_, code) => {176    codeBlocks.push(code);177    return `<<CODE_BLOCK_${codeBlocks.length - 1}>>`;178  });179 180  // Step 2: Protect existing LaTeX expressions181  const latexExpressions: string[] = [];182 183  // Protect block math ($$...$$), \[...\], and \(...\) as before.184  content = content.replace(185    /(\$\$[\s\S]*?\$\$|\\\[[\s\S]*?\\\]|\\\(.*?\\\))/g,186    (match) => {187      latexExpressions.push(match);188      return `<<LATEX_${latexExpressions.length - 1}>>`;189    }190  );191 192  // Protect inline math ($...$) only if it does NOT match a currency pattern.193  // We assume a currency pattern is one where the inner content is purely numeric (with optional decimals).194  content = content.replace(/\$([^$]+)\$/g, (match, inner) => {195    if (/^\s*\d+(?:\.\d+)?\s*$/.test(inner)) {196      // This looks like a currency value (e.g. "$123" or "$12.34"),197      // so don't protect it.198      return match;199    } else {200      // Otherwise, treat it as a LaTeX expression.201      latexExpressions.push(match);202      return `<<LATEX_${latexExpressions.length - 1}>>`;203    }204  });205 206  // Step 3: Escape dollar signs that are likely currency indicators.207  // (Now that inline math is protected, this will only escape dollars not already protected)208  content = content.replace(/\$(?=\d)/g, '\\$');209 210  // Step 4: Restore LaTeX expressions211  content = content.replace(212    /<<LATEX_(\d+)>>/g,213    (_, index) => latexExpressions[parseInt(index)]214  );215 216  // Step 5: Restore code blocks217  content = content.replace(218    /<<CODE_BLOCK_(\d+)>>/g,219    (_, index) => codeBlocks[parseInt(index)]220  );221 222  // Step 6: Apply additional escaping functions223  content = escapeBrackets(content);224  content = escapeMhchem(content);225 226  return content;227}228 229export function escapeBrackets(text: string): string {230  const pattern =231    /(```[\S\s]*?```|`.*?`)|\\\[([\S\s]*?[^\\])\\]|\\\((.*?)\\\)/g;232  return text.replace(233    pattern,234    (235      match: string,236      codeBlock: string | undefined,237      squareBracket: string | undefined,238      roundBracket: string | undefined239    ): string => {240      if (codeBlock != null) {241        return codeBlock;242      } else if (squareBracket != null) {243        return `$$${squareBracket}$$`;244      } else if (roundBracket != null) {245        return `$${roundBracket}$`;246      }247      return match;248    }249  );250}251 252export function escapeMhchem(text: string) {253  return text.replaceAll('$\\ce{', '$\\\\ce{').replaceAll('$\\pu{', '$\\\\pu{');254}255