From 9435d2044a569c1674d4900093ead8126a1e4989 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 14 Aug 2024 21:16:07 +0200
Subject: [PATCH] refac: more robust mermaid chart rendering

---
 .../components/chat/Messages/CodeBlock.svelte | 44 ++++++++++++-------
 .../chat/Messages/MarkdownTokens.svelte       |  2 +-
 src/routes/(app)/+layout.svelte               |  3 ++
 3 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index ebb69a164..bb10af409 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -3,7 +3,7 @@
 	import { loadPyodide } from 'pyodide';
 	import mermaid from 'mermaid';
 
-	import { getContext, getAllContexts } from 'svelte';
+	import { getContext, getAllContexts, onMount } from 'svelte';
 	import { copyToClipboard } from '$lib/utils';
 
 	import 'highlight.js/styles/github-dark.min.css';
@@ -18,6 +18,8 @@
 	export let lang = '';
 	export let code = '';
 
+	let mermaidHtml = null;
+
 	let highlightedCode = null;
 	let executing = false;
 
@@ -213,17 +215,14 @@ __builtins__.input = input`);
 
 	$: if (code) {
 		if (lang === 'mermaid' && (token?.raw ?? '').endsWith('```')) {
-			// Function to perform the code highlighting
-			const renderMermaid = async () => {
-				// mermaid.initialize({ startOnLoad: true });
-				await mermaid.run({
-					querySelector: `.mermaid-${id}`
-				});
-			};
-			// Clear the previous timeout if it exists
-			clearTimeout(debounceTimeout);
-			// Set a new timeout to debounce the code highlighting
-			debounceTimeout = setTimeout(renderMermaid, 50);
+			(async () => {
+				try {
+					const { svg } = await mermaid.render(`mermaid-${id}`, code);
+					mermaidHtml = svg;
+				} catch (error) {
+					console.error('Error:', error);
+				}
+			})();
 		} else {
 			// Function to perform the code highlighting
 			const highlightCode = () => {
@@ -236,13 +235,28 @@ __builtins__.input = input`);
 			debounceTimeout = setTimeout(highlightCode, 10);
 		}
 	}
+
+	onMount(async () => {
+		await mermaid.initialize({ startOnLoad: true });
+
+		if (lang === 'mermaid' && (token?.raw ?? '').endsWith('```')) {
+			try {
+				const { svg } = await mermaid.render(`mermaid-${id}`, code);
+				mermaidHtml = svg;
+			} catch (error) {
+				console.error('Error:', error);
+			}
+		}
+	});
 </script>
 
 <div class="my-2" dir="ltr">
 	{#if lang === 'mermaid'}
-		{#key code}
-			<pre class="mermaid-{id}">{code}</pre>
-		{/key}
+		{#if mermaidHtml}
+			{@html mermaidHtml}
+		{:else}
+			<pre class=" mermaid-{id}">{code}</pre>
+		{/if}
 	{:else}
 		<div
 			class="flex justify-between bg-[#202123] text-white text-xs px-4 pt-1 pb-0.5 rounded-t-lg overflow-x-auto"
diff --git a/src/lib/components/chat/Messages/MarkdownTokens.svelte b/src/lib/components/chat/Messages/MarkdownTokens.svelte
index 8c649180c..028524aa5 100644
--- a/src/lib/components/chat/Messages/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownTokens.svelte
@@ -26,7 +26,7 @@
 		</svelte:element>
 	{:else if token.type === 'code'}
 		<CodeBlock
-			{id}
+			id={`${id}-${tokenIdx}`}
 			{token}
 			lang={token?.lang ?? ''}
 			code={revertSanitizedResponseContent(token?.text ?? '')}
diff --git a/src/routes/(app)/+layout.svelte b/src/routes/(app)/+layout.svelte
index 875ebf4eb..b57b58675 100644
--- a/src/routes/(app)/+layout.svelte
+++ b/src/routes/(app)/+layout.svelte
@@ -3,6 +3,8 @@
 	import { onMount, tick, getContext } from 'svelte';
 	import { openDB, deleteDB } from 'idb';
 	import fileSaver from 'file-saver';
+	import mermaid from 'mermaid';
+
 	const { saveAs } = fileSaver;
 
 	import { goto } from '$app/navigation';
@@ -178,6 +180,7 @@
 			await tick();
 		}
 
+		await mermaid.initialize({ startOnLoad: false });
 		loaded = true;
 	});
 </script>