diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index d2e3a6323d..7e6e72b6bb 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -326,12 +326,26 @@
 	const render = async () => {
 		onUpdate(token);
 		if (lang === 'mermaid' && (token?.raw ?? '').slice(-4).includes('```')) {
-			mermaidHtml = await renderMermaidDiagram(code, $i18n);
+			try {
+				mermaidHtml = await renderMermaidDiagram(code);
+			} catch (error) {
+				console.error('Failed to render mermaid diagram:', error);
+			    const errorMsg = error instanceof Error ? error.message : String(error);
+			    toast.error($i18n.t('Failed to render diagram') + `: ${errorMsg}`);
+			    mermaidHtml = null; 
+			}
 		} else if (
 			(lang === 'vega' || lang === 'vega-lite') &&
 			(token?.raw ?? '').slice(-4).includes('```')
 		) {
-			vegaHtml = await renderVegaVisualization(code, $i18n);
+			try {
+				vegaHtml = await renderVegaVisualization(code);
+			} catch (error) {
+				console.error('Failed to render Vega visualization:', error);
+			    const errorMsg = error instanceof Error ? error.message : String(error);
+			    toast.error($i18n.t('Failed to render diagram') + `: ${errorMsg}`);
+			    vegaHtml = null;
+			}
 		}
 	};
 
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index 856f7fa028..f1e44920a9 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -1581,56 +1581,30 @@ export const decodeString = (str: string) => {
 	}
 };
 
-export const renderMermaidDiagram = async (code: string, i18n?: any) => {
-	try {
-		const { default: mermaid } = await import('mermaid');
-		mermaid.initialize({
-			startOnLoad: false, // Should be false when using render API
-			theme: document.documentElement.classList.contains('dark') ? 'dark' : 'default',
-			securityLevel: 'loose'
-		});
-		const parseResult = await mermaid.parse(code, { suppressErrors: false });
-		if (parseResult) {
-			const { svg } = await mermaid.render(`mermaid-${uuidv4()}`, code);
-			return svg;
-		}
-		return '';
-	} catch (error) {
-		console.error('Failed to render mermaid diagram:', error);
-		const errorMsg = error instanceof Error ? error.message : String(error);
-		toast.error(i18n.t('Failed to render diagram') + `: ${errorMsg}`);
-		return '';
+export const renderMermaidDiagram = async (code: string) => {
+	const { default: mermaid } = await import('mermaid');
+	mermaid.initialize({
+		startOnLoad: false, // Should be false when using render API
+		theme: document.documentElement.classList.contains('dark') ? 'dark' : 'default',
+		securityLevel: 'loose'
+	});
+	const parseResult = await mermaid.parse(code, { suppressErrors: false });
+	if (parseResult) {
+		const { svg } = await mermaid.render(`mermaid-${uuidv4()}`, code);
+		return svg;
 	}
+	return '';
 };
 
 export const renderVegaVisualization = async (spec: string, i18n?: any) => {
-	try {
-		const vega = await import('vega');
-		const parsedSpec = JSON.parse(spec);
-		if (!parsedSpec.$schema) {
-			const errorMsg = 'Specification missing $schema property';
-			toast.error(i18n.t('Failed to render diagram') + `: ${errorMsg}`);
-			throw new Error(errorMsg);
-		}
-		let vegaSpec;
-		if (parsedSpec.$schema.includes('vega-lite')) {
-			const vegaLite = await import('vega-lite');
-			vegaSpec = vegaLite.compile(parsedSpec).spec;
-		} else if (parsedSpec.$schema.includes('vega')) {
-			vegaSpec = parsedSpec;
-		} else {
-			const errorMsg = 'Unknown schema format: ' + parsedSpec.$schema;
-			toast.error(i18n.t('Failed to render diagram') + `: ${errorMsg}`);
-			throw new Error(errorMsg);
-		}
-		const runtime = vega.parse(vegaSpec);
-		const view = new vega.View(runtime, { renderer: 'none' });
-		const svg = await view.toSVG();
-		return svg;
-	} catch (error) {
-		console.error('Failed to render Vega visualization:', error);
-		const errorMsg = error instanceof Error ? error.message : String(error);
-		toast.error(i18n.t('Failed to render diagram') + `: ${errorMsg}`);
-		return '';
+	const vega = await import('vega');
+	const parsedSpec = JSON.parse(spec);
+	let vegaSpec = parsedSpec;
+	if (parsedSpec.$schema && parsedSpec.$schema.includes('vega-lite')) {
+		const vegaLite = await import('vega-lite');
+		vegaSpec = vegaLite.compile(parsedSpec).spec;
 	}
+	const view = new vega.View(vega.parse(vegaSpec), { renderer: 'none' });
+	const svg = await view.toSVG();
+	return svg;
 };