From 92e77d7b33fc8498a3164a5cc30eb64110978020 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 9 Aug 2024 00:01:38 +0200
Subject: [PATCH] refac: katex

---
 package-lock.json                             | 18 +++++
 package.json                                  |  1 +
 .../chat/Messages/KatexRenderer.svelte        |  9 +++
 .../chat/Messages/MarkdownInlineTokens.svelte | 16 +++-
 .../chat/Messages/MarkdownTokens.svelte       | 14 +++-
 .../chat/Messages/ResponseMessage.svelte      | 16 ++--
 src/lib/utils/index.ts                        |  5 +-
 src/lib/utils/katex-extension.ts              | 80 +++++++++++++++++++
 8 files changed, 147 insertions(+), 12 deletions(-)
 create mode 100644 src/lib/components/chat/Messages/KatexRenderer.svelte
 create mode 100644 src/lib/utils/katex-extension.ts

diff --git a/package-lock.json b/package-lock.json
index b7e714ed7..2006a3fd3 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -29,6 +29,7 @@
 				"js-sha256": "^0.10.1",
 				"katex": "^0.16.9",
 				"marked": "^9.1.0",
+				"marked-katex-extension": "^5.1.1",
 				"mermaid": "^10.9.1",
 				"pyodide": "^0.26.1",
 				"socket.io-client": "^4.2.0",
@@ -1544,6 +1545,11 @@
 			"integrity": "sha512-5+fP8P8MFNC+AyZCDxrB2pkZFPGzqQWUzpSeuuVLvm8VMcorNYavBqoFcxK8bQz4Qsbn4oUEEem4wDLfcysGHA==",
 			"dev": true
 		},
+		"node_modules/@types/katex": {
+			"version": "0.16.7",
+			"resolved": "https://registry.npmjs.org/@types/katex/-/katex-0.16.7.tgz",
+			"integrity": "sha512-HMwFiRujE5PjrgwHQ25+bsLJgowjGjm5Z8FVSf0N6PwgJrwxH0QxzHYDcKsTfV3wva0vzrpqMTJS2jXPr5BMEQ=="
+		},
 		"node_modules/@types/mdast": {
 			"version": "3.0.15",
 			"resolved": "https://registry.npmjs.org/@types/mdast/-/mdast-3.0.15.tgz",
@@ -6036,6 +6042,18 @@
 				"node": ">= 16"
 			}
 		},
+		"node_modules/marked-katex-extension": {
+			"version": "5.1.1",
+			"resolved": "https://registry.npmjs.org/marked-katex-extension/-/marked-katex-extension-5.1.1.tgz",
+			"integrity": "sha512-piquiCyZpZ1aiocoJlJkRXr+hkk5UI4xw9GhRZiIAAgvX5rhzUDSJ0seup1JcsgueC8MLNDuqe5cRcAzkFE42Q==",
+			"dependencies": {
+				"@types/katex": "^0.16.7"
+			},
+			"peerDependencies": {
+				"katex": ">=0.16 <0.17",
+				"marked": ">=4 <15"
+			}
+		},
 		"node_modules/matcher-collection": {
 			"version": "2.0.1",
 			"resolved": "https://registry.npmjs.org/matcher-collection/-/matcher-collection-2.0.1.tgz",
diff --git a/package.json b/package.json
index 0c7a8518a..bd2b173bf 100644
--- a/package.json
+++ b/package.json
@@ -70,6 +70,7 @@
 		"js-sha256": "^0.10.1",
 		"katex": "^0.16.9",
 		"marked": "^9.1.0",
+		"marked-katex-extension": "^5.1.1",
 		"mermaid": "^10.9.1",
 		"pyodide": "^0.26.1",
 		"socket.io-client": "^4.2.0",
diff --git a/src/lib/components/chat/Messages/KatexRenderer.svelte b/src/lib/components/chat/Messages/KatexRenderer.svelte
new file mode 100644
index 000000000..766454c29
--- /dev/null
+++ b/src/lib/components/chat/Messages/KatexRenderer.svelte
@@ -0,0 +1,9 @@
+<script lang="ts">
+	import katex from 'katex';
+	import 'katex/contrib/mhchem';
+
+	export let content: string;
+	export let displayMode: boolean = false;
+</script>
+
+{@html katex.renderToString(content, { displayMode, throwOnError: false })}
diff --git a/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte b/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
index 7f984264b..170429f4b 100644
--- a/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
@@ -1,8 +1,11 @@
 <script lang="ts">
 	import type { Token } from 'marked';
-	import { unescapeHtml } from '$lib/utils';
+	import { revertSanitizedResponseContent, unescapeHtml } from '$lib/utils';
 	import { onMount } from 'svelte';
 	import Image from '$lib/components/common/Image.svelte';
+
+	import KatexRenderer from './KatexRenderer.svelte';
+
 	export let id: string;
 	export let tokens: Token[];
 </script>
@@ -25,14 +28,21 @@
 			<svelte:self id={`${id}-em`} tokens={token.tokens} />
 		</em>
 	{:else if token.type === 'codespan'}
-		<code class="codespan">{unescapeHtml(token.text.replaceAll('&amp;', '&'))}</code>
+		<code class="codespan">{revertSanitizedResponseContent(token.raw)}</code>
 	{:else if token.type === 'br'}
 		<br />
 	{:else if token.type === 'del'}
 		<del>
 			<svelte:self id={`${id}-del`} tokens={token.tokens} />
 		</del>
+	{:else if token.type === 'inlineKatex'}
+		{#if token.text}
+			<KatexRenderer
+				content={revertSanitizedResponseContent(token.text)}
+				displayMode={token?.displayMode ?? false}
+			/>
+		{/if}
 	{:else if token.type === 'text'}
-		{unescapeHtml(token.text)}
+		{token.raw}
 	{/if}
 {/each}
diff --git a/src/lib/components/chat/Messages/MarkdownTokens.svelte b/src/lib/components/chat/Messages/MarkdownTokens.svelte
index 1f27fb4eb..089a0847a 100644
--- a/src/lib/components/chat/Messages/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownTokens.svelte
@@ -1,17 +1,22 @@
 <script lang="ts">
+	import { onMount } from 'svelte';
 	import type { Token } from 'marked';
 	import { revertSanitizedResponseContent, unescapeHtml } from '$lib/utils';
+
 	import CodeBlock from '$lib/components/chat/Messages/CodeBlock.svelte';
-	import { onMount } from 'svelte';
 	import MarkdownInlineTokens from '$lib/components/chat/Messages/MarkdownInlineTokens.svelte';
+	import KatexRenderer from './KatexRenderer.svelte';
+
 	export let id: string;
 	export let tokens: Token[];
 	export let top = true;
+
 	const headerComponent = (depth: number) => {
 		return 'h' + depth;
 	};
 </script>
 
+<!-- {JSON.stringify(tokens)} -->
 {#each tokens as token, tokenIdx}
 	{#if token.type === 'hr'}
 		<hr />
@@ -104,6 +109,13 @@
 		{:else}
 			{unescapeHtml(token.text)}
 		{/if}
+	{:else if token.type === 'inlineKatex'}
+		{#if token.text}
+			<KatexRenderer
+				content={revertSanitizedResponseContent(token.text)}
+				displayMode={token?.displayMode ?? false}
+			/>
+		{/if}
 	{:else if token.type === 'space'}
 		{''}
 	{:else}
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index 1226b5231..2ecbd418c 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -4,7 +4,6 @@
 	import { marked } from 'marked';
 	import tippy from 'tippy.js';
 	import auto_render from 'katex/dist/contrib/auto-render.mjs';
-	import 'katex/dist/katex.min.css';
 	import mermaid from 'mermaid';
 
 	import { fade } from 'svelte/transition';
@@ -79,19 +78,24 @@
 
 	let tokens;
 
+	import 'katex/dist/katex.min.css';
+
+	import markedKatex from '$lib/utils/katex-extension';
+
+	const options = {
+		throwOnError: false
+	};
+
+	marked.use(markedKatex(options));
+
 	$: (async () => {
 		if (message?.content) {
 			tokens = marked.lexer(
 				replaceTokens(sanitizeResponseContent(message?.content), model?.name, $user?.name)
 			);
-			// console.log(message?.content, tokens);
 		}
 	})();
 
-	$: if (message) {
-		renderStyling();
-	}
-
 	const renderStyling = async () => {
 		await tick();
 
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index 5b5c69047..de3a2ca2c 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -25,7 +25,8 @@ const convertLatexToSingleLine = (content) => {
 
 export const sanitizeResponseContent = (content: string) => {
 	// replace single backslash with double backslash
-	content = content.replace(/\\/g, '\\\\');
+	content = content.replace(/\\\\/g, '\\\\\\\\');
+
 	content = convertLatexToSingleLine(content);
 
 	// First, temporarily replace valid <video> tags with a placeholder
@@ -87,7 +88,7 @@ export const replaceTokens = (content, char, user) => {
 };
 
 export const revertSanitizedResponseContent = (content: string) => {
-	return content.replaceAll('&lt;', '<').replaceAll('&gt;', '>');
+	return content.replaceAll('&lt;', '<').replaceAll('&gt;', '>').replaceAll('\\\\', '\\');
 };
 
 export function unescapeHtml(html: string) {
diff --git a/src/lib/utils/katex-extension.ts b/src/lib/utils/katex-extension.ts
new file mode 100644
index 000000000..4d0bd99e5
--- /dev/null
+++ b/src/lib/utils/katex-extension.ts
@@ -0,0 +1,80 @@
+import katex from 'katex';
+
+const inlineRule = /^(\${1,2})(?!\$)((?:\\.|[^\\\n])*?(?:\\.|[^\\\n\$]))\1(?=[\s?!\.,:？！。，：]|$)/;
+const inlineRuleNonStandard = /^(\${1,2})(?!\$)((?:\\.|[^\\\n])*?(?:\\.|[^\\\n\$]))\1/; // Non-standard, even if there are no spaces before and after $ or $$, try to parse
+
+const blockRule = /^(\${1,2})\n((?:\\[^]|[^\\])+?)\n\1(?:\n|$)/;
+
+export default function(options = {}) {
+  return {
+    extensions: [
+      inlineKatex(options, createRenderer(options, false)),
+      blockKatex(options, createRenderer(options, true)),
+    ],
+  };
+}
+
+function createRenderer(options, newlineAfter) {
+  return (token) => katex.renderToString(token.text, { ...options, displayMode: token.displayMode }) + (newlineAfter ? '\n' : '');
+}
+
+function inlineKatex(options, renderer) {
+  const nonStandard = options && options.nonStandard;
+  const ruleReg = nonStandard ? inlineRuleNonStandard : inlineRule;
+  return {
+    name: 'inlineKatex',
+    level: 'inline',
+    start(src) {
+      let index;
+      let indexSrc = src;
+
+      while (indexSrc) {
+        index = indexSrc.indexOf('$');
+        if (index === -1) {
+          return;
+        }
+        const f = nonStandard ? index > -1 : index === 0 || indexSrc.charAt(index - 1) === ' ';
+        if (f) {
+          const possibleKatex = indexSrc.substring(index);
+
+          if (possibleKatex.match(ruleReg)) {
+            return index;
+          }
+        }
+
+        indexSrc = indexSrc.substring(index + 1).replace(/^\$+/, '');
+      }
+    },
+    tokenizer(src, tokens) {
+      const match = src.match(ruleReg);
+      if (match) {
+        return {
+          type: 'inlineKatex',
+          raw: match[0],
+          text: match[2].trim(),
+          displayMode: match[1].length === 2,
+        };
+      }
+    },
+    renderer,
+  };
+}
+
+function blockKatex(options, renderer) {
+  return {
+    name: 'blockKatex',
+    level: 'block',
+    tokenizer(src, tokens) {
+      const match = src.match(blockRule);
+      if (match) {
+        return {
+          type: 'blockKatex',
+          raw: match[0],
+          text: match[2].trim(),
+          displayMode: match[1].length === 2,
+        };
+      }
+    },
+    renderer,
+  };
+}
\ No newline at end of file