From 0066760f86e9895e49aa7378c1b5b961b01eed96 Mon Sep 17 00:00:00 2001
From: Morgan <me@morgan.kr>
Date: Tue, 23 Apr 2024 09:10:57 +0900
Subject: [PATCH] Added KaTeX Rendering

---
 .env                       |  1 +
 .gitignore                 |  2 +-
 package-lock.json          | 26 +++++++++++++
 package.json               |  1 +
 src/lib/Code.svelte        | 26 +++++++++++--
 src/lib/Codespan.svelte    | 20 ++++++++++
 src/lib/EditMessage.svelte | 76 ++++++++++++++++++++++++++++++++++++--
 7 files changed, 144 insertions(+), 8 deletions(-)
 create mode 100644 src/lib/Codespan.svelte

diff --git a/.env b/.env
index 38f7106..76bd0df 100644
--- a/.env
+++ b/.env
@@ -2,3 +2,4 @@
 #VITE_API_BASE=http://localhost:5174
 #VITE_ENDPOINT_COMPLETIONS=/v1/chat/completions
 #VITE_ENDPOINT_MODELS=/v1/models
+#VITE_RENDER_LATEX=false
diff --git a/.gitignore b/.gitignore
index d8e5134..77b4273 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,7 +6,7 @@ yarn-debug.log*
 yarn-error.log*
 pnpm-debug.log*
 lerna-debug.log*
-
+yarn.lock
 node_modules
 dist
 dist-ssr
diff --git a/package-lock.json b/package-lock.json
index 6f63e20..6e51f89 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -29,6 +29,7 @@
         "eslint-plugin-svelte3": "^4.0.0",
         "flourite": "^1.2.4",
         "gpt-tokenizer": "^2.1.2",
+        "katex": "^0.16.10",
         "llama-tokenizer-js": "^1.1.3",
         "postcss": "^8.4.32",
         "sass": "^1.69.7",
@@ -3273,6 +3274,31 @@
         "json5": "lib/cli.js"
       }
     },
+    "node_modules/katex": {
+      "version": "0.16.10",
+      "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.10.tgz",
+      "integrity": "sha512-ZiqaC04tp2O5utMsl2TEZTXxa6WSC4yo0fv5ML++D3QZv/vx2Mct0mTlRx3O+uUkjfuAgOkzsCmq5MiUEsDDdA==",
+      "dev": true,
+      "funding": [
+        "https://opencollective.com/katex",
+        "https://github.com/sponsors/katex"
+      ],
+      "dependencies": {
+        "commander": "^8.3.0"
+      },
+      "bin": {
+        "katex": "cli.js"
+      }
+    },
+    "node_modules/katex/node_modules/commander": {
+      "version": "8.3.0",
+      "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz",
+      "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==",
+      "dev": true,
+      "engines": {
+        "node": ">= 12"
+      }
+    },
     "node_modules/keyv": {
       "version": "4.5.4",
       "resolved": "https://registry.npmjs.org/keyv/-/keyv-4.5.4.tgz",
diff --git a/package.json b/package.json
index 84457fa..3080273 100644
--- a/package.json
+++ b/package.json
@@ -35,6 +35,7 @@
     "eslint-plugin-svelte3": "^4.0.0",
     "flourite": "^1.2.4",
     "gpt-tokenizer": "^2.1.2",
+    "katex": "^0.16.10",
     "llama-tokenizer-js": "^1.1.3",
     "postcss": "^8.4.32",
     "sass": "^1.69.7",
diff --git a/src/lib/Code.svelte b/src/lib/Code.svelte
index f570a10..8000981 100644
--- a/src/lib/Code.svelte
+++ b/src/lib/Code.svelte
@@ -27,12 +27,26 @@
     type LanguageType
   } from 'svelte-highlight/languages/index'
 
+  import katex from 'katex'
+  import 'katex/contrib/mhchem'
+
   export const type: 'code' = 'code'
   export const raw: string = ''
   export const codeBlockStyle: 'indented' | undefined = undefined
   export let lang: string | undefined
   export let text: string
 
+  let renderedMath: string | undefined
+
+  $: if (lang === 'rendermath') {
+    renderedMath = katex.renderToString(text, {
+      throwOnError: false,
+      displayMode: true
+    })
+  } else {
+    renderedMath = undefined
+  }
+  
   // Map lang string to LanguageType
   let language: LanguageType<string>
 
@@ -117,7 +131,11 @@
   {@html style}
 </svelte:head>
 
-<div class="code-block is-relative">
-  <button class="button is-light is-outlined is-small p-2" on:click={copyFunction}>Copy</button>
-  <Highlight code={text} {language} />
-</div>
+{#if lang === 'rendermath'}
+  {@html renderedMath}
+{:else}
+  <div class="code-block is-relative">
+    <button class="button is-light is-outlined is-small p-2" on:click={copyFunction}>Copy</button>
+    <Highlight code={text} {language} />
+  </div>
+{/if}
\ No newline at end of file
diff --git a/src/lib/Codespan.svelte b/src/lib/Codespan.svelte
new file mode 100644
index 0000000..2ccf5cd
--- /dev/null
+++ b/src/lib/Codespan.svelte
@@ -0,0 +1,20 @@
+<script lang="ts">
+  export let raw
+  import katex from 'katex'
+  import 'katex/contrib/mhchem'
+
+  let renderedMath: string | undefined
+  if (raw.startsWith('`rendermath')) {
+    renderedMath = katex.renderToString(raw.replace(/`rendermath|`/g, ''), {
+      throwOnError: false,
+      displayMode: false
+    })
+  }
+
+</script>
+
+{#if renderedMath}
+  {@html renderedMath}
+{:else}
+  <code>{raw.replace(/`/g, '')}</code>
+{/if}
diff --git a/src/lib/EditMessage.svelte b/src/lib/EditMessage.svelte
index 596c961..f737be1 100644
--- a/src/lib/EditMessage.svelte
+++ b/src/lib/EditMessage.svelte
@@ -1,18 +1,21 @@
 <script lang="ts">
   import Code from './Code.svelte'
+  import Codespan from './Codespan.svelte'
   import { afterUpdate, createEventDispatcher, onMount } from 'svelte'
   import { deleteMessage, deleteSummaryMessage, truncateFromMessage, submitExitingPromptsNow, continueMessage, updateMessages } from './Storage.svelte'
   import { getPrice } from './Stats.svelte'
   import SvelteMarkdown from 'svelte-markdown'
   import type { Message, Model, Chat } from './Types.svelte'
   import Fa from 'svelte-fa/src/fa.svelte'
-  import { faTrash, faDiagramPredecessor, faDiagramNext, faCircleCheck, faPaperPlane, faEye, faEyeSlash, faEllipsis, faDownload, faClipboard } from '@fortawesome/free-solid-svg-icons/index'
+  import { faTrash, faDiagramPredecessor, faDiagramNext, faCircleCheck, faPaperPlane, faEye, faEyeSlash, faEllipsis, faDownload, faClipboard, faSquareRootVariable } from '@fortawesome/free-solid-svg-icons/index'
   import { errorNotice, scrollToMessage } from './Util.svelte'
   import { openModal } from 'svelte-modals'
   import PromptConfirm from './PromptConfirm.svelte'
   import { getImage } from './ImageStore.svelte'
   import { getModelDetail } from './Models.svelte'
 
+  import 'katex/dist/katex.min.css'
+
   export let message:Message
   export let chatId:number
   export let chat:Chat
@@ -31,6 +34,12 @@
     breaks: true, // Enable line breaks in markdown
     mangle: false // Do not mangle email addresses
   }
+  
+  const renderers = {
+    code: Code,
+    html: Code,
+    codespan: Codespan
+  }
 
   const getDisplayMessage = ():string => {
     const content = message.content
@@ -214,6 +223,57 @@
     document.body.removeChild(a)
   }
 
+  const preprocessMath = (text: string): string => {
+    let codeBlockPlaceholderPrefix = '__prefix__c0d3b10ck__'
+    while (text.indexOf(codeBlockPlaceholderPrefix) > 0) {
+      codeBlockPlaceholderPrefix = codeBlockPlaceholderPrefix + '_'
+    }
+    let index = 0
+    const codeBlocks = []
+
+    const codeBlockRegex = /(```[\s\S]*?```|`[^`]*`)/g
+  
+    text = text.replace(codeBlockRegex, (match) => {
+      const placeholder = `${codeBlockPlaceholderPrefix}idx${index}__`
+      codeBlocks.push(match)
+      index++
+      return placeholder
+    })
+
+    text = text
+      .replace(/(\\\[((?:\s|\S)*?)\\\])|(\$\$((?:\s|\S)*?)\$\$)/g, (match, p1, p2, p3, p4) => {
+        const math = p2 || p4
+        return '\n```rendermath\n' + math.trim() + '\n```\n'
+      })
+      .replace(/(\\\((?!\$)(.*?)\\\))|(?<!\\|\$)\$(?!\$)(.*?[^\\])\$(?!\$)/g, (match, p1, p2, p3) => {
+        const math = p2 || p3
+        return '`rendermath' + math.trim() + '`'
+      })
+
+    // .replace(/\\\[((?:\s|\S)*?)\\\]/g, (match, math) => {
+    //     return '\n```rendermath\n' + math.trim() + '\n```\n'
+    // })
+    // .replace(/\$\$((?:\s|\S)*?)\$\$/g, (match, math) => {
+    //     return '\n```rendermath\n' + math.trim() + '\n```\n'
+    // })
+    // .replace(/\\\((?!\$)(.*?[^\\])\\\)/g, (match, math) => {
+    //     return '`rendermath' + math.trim() + '`'
+    // })
+    // .replace(/(?<!\\|\$)\$(?!\$)(.*?[^\\])\$(?!\$)/g, (match, math) => {
+    //     return '`rendermath' + math.trim() + '`'
+    // })
+
+    text = text.replace(new RegExp(`${codeBlockPlaceholderPrefix}idx(\\d+)__`, 'g'), (match, p1) => {
+      return codeBlocks[p1]
+    })
+
+    return text
+  }
+
+  const renderMathMsg = () => {
+    displayMessage = preprocessMath(message.content);
+  };
+
 </script>
 
 <article
@@ -253,9 +313,9 @@
         {/if}
         {#key refreshCounter}
         <SvelteMarkdown 
-          source={displayMessage} 
+          bind:source={displayMessage} 
           options={markdownOptions} 
-          renderers={{ code: Code, html: Code }}
+          renderers={renderers}
         />
         {/key}
         {#if imageUrl}
@@ -371,6 +431,16 @@
         <span class="icon"><Fa icon={faClipboard} /></span>
         </a>
       {/if}
+      <a
+        href={'#'}
+        title="Render LaTeX in message"
+        class="button is-small"
+        on:click|preventDefault={() => {
+          renderMathMsg()
+        }}
+      >
+      <span class="icon"><Fa icon={faSquareRootVariable} /></span>
+      </a>
       {#if imageUrl}
         <a
           href={'#'}