From 003968f06a260ab1ac7418c2d30bd85ed4ba2747 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Tue, 18 Feb 2025 20:15:16 -0800
Subject: [PATCH] enh: jupyter matplotlib support

---
 .../components/chat/Messages/CodeBlock.svelte | 99 ++++++++++++++++++-
 1 file changed, 94 insertions(+), 5 deletions(-)

diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index c4614d52b..59f58d2e3 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -130,9 +130,65 @@
 			});
 
 			if (output) {
-				stdout = output.stdout;
-				stderr = output.stderr;
-				result = output.result;
+				if (output['stdout']) {
+					stdout = output['stdout'];
+					const stdoutLines = stdout.split('\n');
+
+					for (const [idx, line] of stdoutLines.entries()) {
+						if (line.startsWith('data:image/png;base64')) {
+							if (files) {
+								files.push({
+									type: 'image/png',
+									data: line
+								});
+							} else {
+								files = [
+									{
+										type: 'image/png',
+										data: line
+									}
+								];
+							}
+
+							if (stdout.startsWith(`${line}\n`)) {
+								stdout = stdout.replace(`${line}\n`, ``);
+							} else if (stdout.startsWith(`${line}`)) {
+								stdout = stdout.replace(`${line}`, ``);
+							}
+						}
+					}
+				}
+
+				if (output['result']) {
+					result = output['result'];
+					const resultLines = result.split('\n');
+
+					for (const [idx, line] of resultLines.entries()) {
+						if (line.startsWith('data:image/png;base64')) {
+							if (files) {
+								files.push({
+									type: 'image/png',
+									data: line
+								});
+							} else {
+								files = [
+									{
+										type: 'image/png',
+										data: line
+									}
+								];
+							}
+
+							if (result.startsWith(`${line}\n`)) {
+								result = result.replace(`${line}\n`, ``);
+							} else if (result.startsWith(`${line}`)) {
+								result = result.replace(`${line}`, ``);
+							}
+						}
+					}
+				}
+
+				output['stderr'] && (stderr = output['stderr']);
 			}
 		} else {
 			executePythonAsWorker(code);
@@ -205,7 +261,40 @@
 							];
 						}
 
-						stdout = stdout.replace(`${line}\n`, ``);
+						if (stdout.startsWith(`${line}\n`)) {
+							stdout = stdout.replace(`${line}\n`, ``);
+						} else if (stdout.startsWith(`${line}`)) {
+							stdout = stdout.replace(`${line}`, ``);
+						}
+					}
+				}
+			}
+
+			if (data['result']) {
+				result = data['result'];
+				const resultLines = result.split('\n');
+
+				for (const [idx, line] of resultLines.entries()) {
+					if (line.startsWith('data:image/png;base64')) {
+						if (files) {
+							files.push({
+								type: 'image/png',
+								data: line
+							});
+						} else {
+							files = [
+								{
+									type: 'image/png',
+									data: line
+								}
+							];
+						}
+
+						if (result.startsWith(`${line}\n`)) {
+							result = result.replace(`${line}\n`, ``);
+						} else if (result.startsWith(`${line}`)) {
+							result = result.replace(`${line}`, ``);
+						}
 					}
 				}
 			}
@@ -391,7 +480,7 @@
 				class="bg-gray-50 dark:bg-[#202123] dark:text-white max-w-full overflow-x-auto scrollbar-hidden"
 			/>
 
-			{#if executing || stdout || stderr || result}
+			{#if executing || stdout || stderr || result || files}
 				<div
 					class="bg-gray-50 dark:bg-[#202123] dark:text-white rounded-b-lg! py-4 px-4 flex flex-col gap-2"
 				>