diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index 33b0572f41..ebb69a1648 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -1,6 +1,8 @@
 <script lang="ts">
 	import hljs from 'highlight.js';
 	import { loadPyodide } from 'pyodide';
+	import mermaid from 'mermaid';
+
 	import { getContext, getAllContexts } from 'svelte';
 	import { copyToClipboard } from '$lib/utils';
 
@@ -12,6 +14,7 @@
 
 	export let id = '';
 
+	export let token;
 	export let lang = '';
 	export let code = '';
 
@@ -207,23 +210,39 @@ __builtins__.input = input`);
 	};
 
 	let debounceTimeout;
+
 	$: if (code) {
-		// Function to perform the code highlighting
-		const highlightCode = () => {
-			highlightedCode = hljs.highlightAuto(code, hljs.getLanguage(lang)?.aliases).value || 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);
+		} else {
+			// Function to perform the code highlighting
+			const highlightCode = () => {
+				highlightedCode = hljs.highlightAuto(code, hljs.getLanguage(lang)?.aliases).value || code;
+			};
 
-		// Clear the previous timeout if it exists
-		clearTimeout(debounceTimeout);
-
-		// Set a new timeout to debounce the code highlighting
-		debounceTimeout = setTimeout(highlightCode, 10);
+			// Clear the previous timeout if it exists
+			clearTimeout(debounceTimeout);
+			// Set a new timeout to debounce the code highlighting
+			debounceTimeout = setTimeout(highlightCode, 10);
+		}
 	}
 </script>
 
 <div class="my-2" dir="ltr">
 	{#if lang === 'mermaid'}
-		<pre class="mermaid">{code}</pre>
+		{#key code}
+			<pre class="mermaid-{id}">{code}</pre>
+		{/key}
 	{: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 5d2c938f68..8c649180ce 100644
--- a/src/lib/components/chat/Messages/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownTokens.svelte
@@ -27,6 +27,7 @@
 	{:else if token.type === 'code'}
 		<CodeBlock
 			{id}
+			{token}
 			lang={token?.lang ?? ''}
 			code={revertSanitizedResponseContent(token?.text ?? '')}
 		/>
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index 9a6c1c8733..8dd332d25d 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -2,7 +2,6 @@
 	import { toast } from 'svelte-sonner';
 	import dayjs from 'dayjs';
 	import { marked } from 'marked';
-	import mermaid from 'mermaid';
 
 	import { fade } from 'svelte/transition';
 	import { createEventDispatcher } from 'svelte';
@@ -277,18 +276,11 @@
 	$: if (!edit) {
 		(async () => {
 			await tick();
-
-			await mermaid.run({
-				querySelector: '.mermaid'
-			});
 		})();
 	}
 
 	onMount(async () => {
 		await tick();
-		await mermaid.run({
-			querySelector: '.mermaid'
-		});
 	});
 </script>