From 516051304e1b1f250c34438746ade673a79bd40c Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 28 Jun 2026 22:33:59 -0500 Subject: [PATCH] refac --- backend/open_webui/main.py | 12 + src/lib/pyodide/createPyodideWorker.ts | 10 + src/lib/pyodide/pyodideKernel.ts | 81 ------- src/lib/pyodide/pyodideKernel.worker.ts | 123 ----------- src/lib/pyodide/pyodideSandboxHost.ts | 280 ++++++++++++++++++++++++ src/routes/+layout.svelte | 4 +- 6 files changed, 304 insertions(+), 206 deletions(-) create mode 100644 src/lib/pyodide/createPyodideWorker.ts delete mode 100644 src/lib/pyodide/pyodideKernel.ts delete mode 100644 src/lib/pyodide/pyodideKernel.worker.ts create mode 100644 src/lib/pyodide/pyodideSandboxHost.ts diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index c2ea70a0a5..795a163915 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -92,6 +92,7 @@ from open_webui.env import ( ENABLE_SCIM, ENABLE_SIGNUP_PASSWORD_CONFIRMATION, ENABLE_STAR_SESSIONS_MIDDLEWARE, + ENABLE_PYODIDE_FILE_PERSISTENCE, ENABLE_VERSION_UPDATE_CHECK, ENABLE_WEBSOCKET_SUPPORT, GLOBAL_LOG_LEVEL, @@ -271,6 +272,13 @@ class SPAStaticFiles(StaticFiles): raise ex +class CORSStaticFiles(StaticFiles): + async def get_response(self, path: str, scope): + response = await super().get_response(path, scope) + response.headers['Access-Control-Allow-Origin'] = '*' + return response + + if LOG_FORMAT != 'json': banner = rf""" ██████╗ ██████╗ ███████╗███╗ ██╗ ██╗ ██╗███████╗██████╗ ██╗ ██╗██╗ @@ -2580,6 +2588,10 @@ applications.get_swagger_ui_html = swagger_ui_html if os.path.exists(FRONTEND_BUILD_DIR): mimetypes.add_type('text/javascript', '.js') + pyodide_dir = FRONTEND_BUILD_DIR / 'pyodide' + if os.path.exists(pyodide_dir): + app.mount('/pyodide', CORSStaticFiles(directory=pyodide_dir), name='pyodide') + app.mount( '/', SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True), diff --git a/src/lib/pyodide/createPyodideWorker.ts b/src/lib/pyodide/createPyodideWorker.ts new file mode 100644 index 0000000000..de711c424c --- /dev/null +++ b/src/lib/pyodide/createPyodideWorker.ts @@ -0,0 +1,10 @@ +import { get } from 'svelte/store'; + +import { config } from '$lib/stores'; +import { PyodideSandboxHost } from '$lib/pyodide/pyodideSandboxHost'; +import PyodideWorker from '$lib/workers/pyodide.worker?worker'; + +export const createPyodideWorker = (): Worker => + get(config)?.features?.enable_pyodide_file_persistence + ? new PyodideWorker() + : (new PyodideSandboxHost() as unknown as Worker); diff --git a/src/lib/pyodide/pyodideKernel.ts b/src/lib/pyodide/pyodideKernel.ts deleted file mode 100644 index bd3eaeb77f..0000000000 --- a/src/lib/pyodide/pyodideKernel.ts +++ /dev/null @@ -1,81 +0,0 @@ -import PyodideWorker from '$lib/pyodide/pyodideKernel.worker?worker'; - -export type CellState = { - id: string; - status: 'idle' | 'running' | 'completed' | 'error'; - result: any; - stdout: string; - stderr: string; -}; - -export class PyodideKernel { - private worker: Worker; - private listeners: Map void>; - - constructor() { - this.worker = new PyodideWorker(); - this.listeners = new Map(); - - // Listen to messages from the worker - this.worker.onmessage = (event) => { - const { type, id, ...data } = event.data; - - if ((type === 'stdout' || type === 'stderr') && this.listeners.has(id)) { - this.listeners.get(id)?.({ type, id, ...data }); - } else if (type === 'result' && this.listeners.has(id)) { - this.listeners.get(id)?.({ type, id, ...data }); - // Remove the listener once the result is delivered - this.listeners.delete(id); - } else if (type === 'kernelState') { - this.listeners.forEach((listener) => listener({ type, ...data })); - } - }; - - // Initialize the worker - this.worker.postMessage({ type: 'initialize' }); - } - - async execute(id: string, code: string): Promise { - return new Promise((resolve, reject) => { - // Set up the listener for streaming and execution result - const state: CellState = { - id, - status: 'running', - result: null, - stdout: '', - stderr: '' - }; - - this.listeners.set(id, (data) => { - if (data.type === 'stdout') { - state.stdout += data.message; - } else if (data.type === 'stderr') { - state.stderr += data.message; - } else if (data.type === 'result') { - // Final result - const { state: finalState } = data; - resolve(finalState); - } - }); - - // Send execute request to the worker - this.worker.postMessage({ type: 'execute', id, code }); - }); - } - - async getState() { - return new Promise>((resolve) => { - this.worker.postMessage({ type: 'getState' }); - this.listeners.set('kernelState', (data) => { - if (data.type === 'kernelState') { - resolve(data.state); - } - }); - }); - } - - terminate() { - this.worker.postMessage({ type: 'terminate' }); - this.worker.terminate(); - } -} diff --git a/src/lib/pyodide/pyodideKernel.worker.ts b/src/lib/pyodide/pyodideKernel.worker.ts deleted file mode 100644 index 4f1eea6da0..0000000000 --- a/src/lib/pyodide/pyodideKernel.worker.ts +++ /dev/null @@ -1,123 +0,0 @@ -import { loadPyodide, type PyodideInterface } from 'pyodide'; - -declare global { - interface Window { - stdout: string | null; - stderr: string | null; - pyodide: PyodideInterface; - cells: Record; - indexURL: string; - } -} - -type CellState = { - id: string; - status: 'idle' | 'running' | 'completed' | 'error'; - result: any; - stdout: string; - stderr: string; -}; - -const initializePyodide = async () => { - // Ensure Pyodide is loaded once and cached in the worker's global scope - if (!self.pyodide) { - self.indexURL = '/pyodide/'; - self.stdout = ''; - self.stderr = ''; - self.cells = {}; - - self.pyodide = await loadPyodide({ - indexURL: self.indexURL - }); - } -}; - -const executeCode = async (id: string, code: string) => { - if (!self.pyodide) { - await initializePyodide(); - } - - // Update the cell state to "running" - self.cells[id] = { - id, - status: 'running', - result: null, - stdout: '', - stderr: '' - }; - - // Redirect stdout/stderr to stream updates - self.pyodide.setStdout({ - batched: (msg: string) => { - self.cells[id].stdout += msg; - self.postMessage({ type: 'stdout', id, message: msg }); - } - }); - self.pyodide.setStderr({ - batched: (msg: string) => { - self.cells[id].stderr += msg; - self.postMessage({ type: 'stderr', id, message: msg }); - } - }); - - try { - // Dynamically load required packages based on imports in the Python code - await self.pyodide.loadPackagesFromImports(code, { - messageCallback: (msg: string) => { - self.postMessage({ type: 'stdout', id, package: true, message: `[package] ${msg}` }); - }, - errorCallback: (msg: string) => { - self.postMessage({ type: 'stderr', id, package: true, message: `[package] ${msg}` }); - } - }); - - // Execute the Python code - const result = await self.pyodide.runPythonAsync(code); - self.cells[id].result = result; - self.cells[id].status = 'completed'; - } catch (error) { - self.cells[id].status = 'error'; - self.cells[id].stderr += `\n${error.toString()}`; - } finally { - // Notify parent thread when execution completes - self.postMessage({ - type: 'result', - id, - state: self.cells[id] - }); - } -}; - -// Handle messages from the main thread -self.onmessage = async (event) => { - const { type, id, code, ...args } = event.data; - - switch (type) { - case 'initialize': - await initializePyodide(); - self.postMessage({ type: 'initialized' }); - break; - - case 'execute': - if (id && code) { - await executeCode(id, code); - } - break; - - case 'getState': - self.postMessage({ - type: 'kernelState', - state: self.cells - }); - break; - - case 'terminate': - // Explicitly clear the worker for cleanup - for (const key in self.cells) delete self.cells[key]; - self.close(); - break; - - default: - console.error(`Unknown message type: ${type}`); - } -}; diff --git a/src/lib/pyodide/pyodideSandboxHost.ts b/src/lib/pyodide/pyodideSandboxHost.ts new file mode 100644 index 0000000000..33c9c6451a --- /dev/null +++ b/src/lib/pyodide/pyodideSandboxHost.ts @@ -0,0 +1,280 @@ +type MessageListener = (event: MessageEvent) => void; +type ErrorListener = (event: Event) => void; +type QueuedMessage = { message: unknown; transfer: Transferable[] }; + +const sandboxScript = String.raw` +(function () { + let pyodide = null; + let pyodideReady = null; + let stdout = null; + let stderr = null; + + function post(message, transfer) { + parent.postMessage(message, '*', transfer || []); + } + + async function loadRuntime(packages) { + stdout = null; + stderr = null; + pyodide = await loadPyodide({ + indexURL: '/pyodide/', + stdout: function (text) { + stdout = stdout ? stdout + text + '\n' : text + '\n'; + }, + stderr: function (text) { + stderr = stderr ? stderr + text + '\n' : text + '\n'; + }, + packages: ['micropip'] + }); + pyodide.FS.mkdirTree('/mnt/uploads'); + await pyodide.pyimport('micropip').install(packages || []); + } + + async function ensureRuntime(packages) { + if (!pyodideReady) pyodideReady = loadRuntime(packages || []); + await pyodideReady; + if (packages && packages.length > 0) { + await pyodide.pyimport('micropip').install(packages); + } + } + + function ensureDir(dir) { + try { + pyodide.FS.stat(dir); + } catch { + pyodide.FS.mkdirTree(dir); + } + } + + function upload(files, dir) { + dir = dir || '/mnt/uploads'; + ensureDir(dir); + for (const file of files || []) { + pyodide.FS.writeFile(dir + '/' + file.name, new Uint8Array(file.data)); + } + } + + function list(path) { + const entries = []; + try { + const names = pyodide.FS.readdir(path).filter(function (name) { + return name !== '.' && name !== '..'; + }); + for (const name of names) { + try { + const stat = pyodide.FS.stat(path + '/' + name); + const isDir = pyodide.FS.isDir(stat.mode); + entries.push({ name: name, type: isDir ? 'directory' : 'file', size: isDir ? 0 : stat.size }); + } catch {} + } + } catch {} + return entries; + } + + function remove(path) { + try { + const stat = pyodide.FS.stat(path); + if (!pyodide.FS.isDir(stat.mode)) { + pyodide.FS.unlink(path); + return; + } + const names = pyodide.FS.readdir(path).filter(function (name) { + return name !== '.' && name !== '..'; + }); + for (const name of names) remove(path + '/' + name); + pyodide.FS.rmdir(path); + } catch {} + } + + function clean(value) { + try { + if (value == null) return null; + if (['string', 'number', 'boolean'].includes(typeof value)) return value; + if (typeof value === 'bigint') return value.toString(); + if (Array.isArray(value)) return value.map(clean); + if (typeof value.toJs === 'function') return clean(value.toJs()); + if (typeof value === 'object') { + const out = {}; + for (const key in value) { + if (Object.prototype.hasOwnProperty.call(value, key)) out[key] = clean(value[key]); + } + return out; + } + return JSON.stringify(value); + } catch (error) { + return '[processResult error]: ' + (error && error.message ? error.message : String(error)); + } + } + + async function patchMatplotlib() { + await pyodide.runPythonAsync([ + 'import base64', + 'import os', + 'from io import BytesIO', + 'os.environ["MPLBACKEND"] = "AGG"', + 'import matplotlib.pyplot', + '_old_show = matplotlib.pyplot.show', + 'assert _old_show, "matplotlib.pyplot.show"', + 'def show(*, block=None):', + '\\tbuf = BytesIO()', + '\\tmatplotlib.pyplot.savefig(buf, format="png")', + '\\tbuf.seek(0)', + '\\timg_str = base64.b64encode(buf.read()).decode("utf-8")', + '\\tmatplotlib.pyplot.clf()', + '\\tbuf.close()', + '\\tprint(f"data:image/png;base64,{img_str}")', + 'matplotlib.pyplot.show = show' + ].join('\n')); + } + + async function execute(id, code, files) { + stdout = null; + stderr = null; + let result = null; + if (files && files.length > 0) upload(files); + try { + if (code.includes('matplotlib')) await patchMatplotlib(); + result = clean(await pyodide.runPythonAsync(code)); + } catch (error) { + stderr = error && error.message ? error.message : String(error); + } + post({ id: id, result: result, stdout: stdout, stderr: stderr }); + } + + window.addEventListener('message', async function (event) { + if (event.source !== parent) return; + const data = event.data || {}; + const id = data.id; + if (!data.type || data.type === 'execute') { + await ensureRuntime(data.packages || []); + await execute(id, data.code, data.files); + return; + } + await ensureRuntime(); + switch (data.type) { + case 'fs:upload': + upload(data.files, data.dir); + post({ id: id, type: data.type, success: true }); + break; + case 'fs:list': + post({ id: id, type: data.type, entries: list(data.path) }); + break; + case 'fs:read': + try { + const buffer = pyodide.FS.readFile(data.path).buffer; + post({ id: id, type: data.type, data: buffer }, [buffer]); + } catch (error) { + post({ id: id, type: data.type, error: error && error.message ? error.message : String(error) }); + } + break; + case 'fs:delete': + remove(data.path); + post({ id: id, type: data.type, success: true }); + break; + case 'fs:mkdir': + pyodide.FS.mkdirTree(data.path); + post({ id: id, type: data.type, success: true }); + break; + case 'fs:sync': + post({ id: id, type: data.type, success: true }); + break; + } + }); +})(); +`; + +const sandboxHtml = ``; + +export class PyodideSandboxHost { + onmessage: MessageListener | null = null; + onerror: ErrorListener | null = null; + + private iframe: HTMLIFrameElement; + private ready = false; + private queue: QueuedMessage[] = []; + private messageListeners = new Set(); + private errorListeners = new Set(); + private onWindowMessage: (event: MessageEvent) => void; + private onIframeLoad: () => void; + private onIframeError: (event: Event) => void; + + constructor() { + this.iframe = document.createElement('iframe'); + this.iframe.setAttribute('sandbox', 'allow-scripts'); + this.iframe.setAttribute('aria-hidden', 'true'); + this.iframe.setAttribute('title', 'pyodide-sandbox'); + this.iframe.style.display = 'none'; + this.iframe.srcdoc = sandboxHtml; + + this.onWindowMessage = (event: MessageEvent) => { + if (event.source !== this.iframe.contentWindow) { + return; + } + + const messageEvent = { data: event.data } as MessageEvent; + this.onmessage?.(messageEvent); + for (const listener of this.messageListeners) { + listener(messageEvent); + } + }; + + this.onIframeLoad = () => { + this.ready = true; + for (const item of this.queue) { + this.post(item.message, item.transfer); + } + this.queue = []; + }; + + this.onIframeError = (event: Event) => { + this.onerror?.(event); + for (const listener of this.errorListeners) { + listener(event); + } + }; + + window.addEventListener('message', this.onWindowMessage); + this.iframe.addEventListener('load', this.onIframeLoad, { once: true }); + this.iframe.addEventListener('error', this.onIframeError); + document.body.appendChild(this.iframe); + } + + postMessage(message: unknown, transfer: Transferable[] = []) { + if (this.ready) { + this.post(message, transfer); + } else { + this.queue.push({ message, transfer }); + } + } + + addEventListener(type: 'message' | 'error', listener: MessageListener | ErrorListener) { + if (type === 'message') { + this.messageListeners.add(listener as MessageListener); + } else if (type === 'error') { + this.errorListeners.add(listener as ErrorListener); + } + } + + removeEventListener(type: 'message' | 'error', listener: MessageListener | ErrorListener) { + if (type === 'message') { + this.messageListeners.delete(listener as MessageListener); + } else if (type === 'error') { + this.errorListeners.delete(listener as ErrorListener); + } + } + + terminate() { + window.removeEventListener('message', this.onWindowMessage); + this.iframe.removeEventListener('load', this.onIframeLoad); + this.iframe.removeEventListener('error', this.onIframeError); + this.messageListeners.clear(); + this.errorListeners.clear(); + this.onmessage = null; + this.onerror = null; + this.iframe.remove(); + } + + private post(message: unknown, transfer: Transferable[]) { + this.iframe.contentWindow?.postMessage(message, '*', transfer); + } +} diff --git a/src/routes/+layout.svelte b/src/routes/+layout.svelte index 5baa0195ec..8a831cee73 100644 --- a/src/routes/+layout.svelte +++ b/src/routes/+layout.svelte @@ -1,7 +1,7 @@