refac: code interpreter

src/lib/components/chat/Messages/CodeBlock.svelte

@@ -50,6 +50,7 @@
 	let stdout = null;
 	let stderr = null;
 	let result = null;
+	let files = null;
 
 	let copied = false;
 	let saved = false;
@@ -110,7 +111,7 @@
 	};
 
 	const executePython = async (code) => {
-		if (!code.includes('input') && !code.includes('matplotlib')) {
+		if (!code.includes('input')) {
 			executePythonAsWorker(code);
 		} else {
 			result = null;
@@ -211,7 +212,8 @@ __builtins__.input = input`);
 			code.includes('re') ? 'regex' : null,
 			code.includes('seaborn') ? 'seaborn' : null,
 			code.includes('sympy') ? 'sympy' : null,
-			code.includes('tiktoken') ? 'tiktoken' : null
+			code.includes('tiktoken') ? 'tiktoken' : null,
+			code.includes('matplotlib') ? 'matplotlib' : null
 		].filter(Boolean);
 
 		console.log(packages);
@@ -238,7 +240,31 @@ __builtins__.input = input`);
 
 			console.log(id, data);
 
-			data['stdout'] && (stdout = data['stdout']);
+			if (data['stdout']) {
+				stdout = data['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
+								}
+							];
+						}
+
+						stdout = stdout.replace(`${line}\n`, ``);
+					}
+				}
+			}
+
 			data['stderr'] && (stderr = data['stderr']);
 			data['result'] && (result = data['result']);
 
@@ -430,10 +456,21 @@ __builtins__.input = input`);
 								<div class="text-sm">{stdout || stderr}</div>
 							</div>
 						{/if}
-						{#if result}
+						{#if result || files}
 							<div class=" ">
 								<div class=" text-gray-500 text-xs mb-1">RESULT</div>
-								<div class="text-sm">{`${JSON.stringify(result)}`}</div>
+								{#if result}
+									<div class="text-sm">{`${JSON.stringify(result)}`}</div>
+								{/if}
+								{#if files}
+									<div class="flex flex-col gap-2">
+										{#each files as file}
+											{#if file.type.startsWith('image')}
+												<img src={file.data} alt="Output" />
+											{/if}
+										{/each}
+									</div>
+								{/if}
 							</div>
 						{/if}
 					{/if}

src/lib/workers/pyodide.worker.ts

@@ -77,6 +77,35 @@ self.onmessage = async (event) => {
 	await loadPyodideAndPackages(self.packages);
 
 	try {
+		// check if matplotlib is imported in the code
+		if (code.includes('matplotlib')) {
+			// Override plt.show() to return base64 image
+			await self.pyodide.runPythonAsync(`import base64
+import os
+from io import BytesIO
+
+# before importing matplotlib
+# to avoid the wasm backend (which needs js.document', not available in worker)
+os.environ["MPLBACKEND"] = "AGG"
+
+import matplotlib.pyplot
+
+_old_show = matplotlib.pyplot.show
+assert _old_show, "matplotlib.pyplot.show"
+
+def show(*, block=None):
+	buf = BytesIO()
+	matplotlib.pyplot.savefig(buf, format="png")
+	buf.seek(0)
+	# encode to a base64 str
+	img_str = base64.b64encode(buf.read()).decode('utf-8')
+	matplotlib.pyplot.clf()
+	buf.close()
+	print(f"data:image/png;base64,{img_str}")
+
+matplotlib.pyplot.show = show`);
+		}
+
 		self.result = await self.pyodide.runPythonAsync(code);
 
 		// Safely process and recursively serialize the result