From 58d2fd9ac15d47c5484467e9bd7bb89e4c774ebe Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Mon, 3 Feb 2025 00:03:41 -0800
Subject: [PATCH] refac: styling

---
 backend/open_webui/utils/middleware.py        | 15 ++++++---
 .../components/chat/Messages/CodeBlock.svelte | 31 +++++++++++++++++++
 .../Messages/Markdown/MarkdownTokens.svelte   |  8 ++++-
 src/lib/components/common/Collapsible.svelte  |  6 ++++
 src/lib/workers/pyodide.worker.ts             |  6 +++-
 5 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index f39066a29..c93c840bf 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -7,6 +7,7 @@ from aiocache import cached
 from typing import Any, Optional
 import random
 import json
+import html
 import inspect
 import re
 
@@ -1082,7 +1083,7 @@ async def process_chat_response(
 
         # Handle as a background task
         async def post_response_handler(response, events):
-            def serialize_content_blocks(content_blocks):
+            def serialize_content_blocks(content_blocks, raw=False):
                 content = ""
 
                 for block in content_blocks:
@@ -1103,12 +1104,16 @@ async def process_chat_response(
 
                     elif block["type"] == "code_interpreter":
                         attributes = block.get("attributes", {})
-                        output = block.get("output", {})
-
+                        output = block.get("output", None)
                         lang = attributes.get("lang", "")
 
                         if output:
-                            content = f'{content}<details type="code_interpreter" done="true">\n<summary>Analyzed</summary>\n```{lang}\n{block["content"]}\n```\n```output\n{output}\n```\n</details>\n'
+                            output = html.escape(json.dumps(output))
+
+                            if raw:
+                                content = f'{content}<details type="code_interpreter" done="true" output="{output}">\n<summary>Analyzed</summary>\n```{lang}\n{block["content"]}\n```\n```output\n{output}\n```\n</details>\n'
+                            else:
+                                content = f'{content}<details type="code_interpreter" done="true" output="{output}">\n<summary>Analyzed</summary>\n```{lang}\n{block["content"]}\n```\n</details>\n'
                         else:
                             content = f'{content}<details type="code_interpreter" done="false">\n<summary>Analyzing...</summary>\n```{lang}\n{block["content"]}\n```\n</details>\n'
 
@@ -1388,7 +1393,7 @@ async def process_chat_response(
                                     {
                                         "role": "assistant",
                                         "content": serialize_content_blocks(
-                                            content_blocks
+                                            content_blocks, raw=True
                                         ),
                                     },
                                 ],
diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte
index 202e4f458..ed95c607e 100644
--- a/src/lib/components/chat/Messages/CodeBlock.svelte
+++ b/src/lib/components/chat/Messages/CodeBlock.svelte
@@ -25,6 +25,7 @@
 	export let token;
 	export let lang = '';
 	export let code = '';
+	export let attributes = {};
 
 	export let className = 'my-2';
 	export let editorClassName = '';
@@ -279,6 +280,36 @@ __builtins__.input = input`);
 
 	$: dispatch('code', { lang, code });
 
+	$: if (attributes) {
+		onAttributesUpdate();
+	}
+
+	const onAttributesUpdate = () => {
+		if (attributes?.output) {
+			// Create a helper function to unescape HTML entities
+			const unescapeHtml = (html) => {
+				const textArea = document.createElement('textarea');
+				textArea.innerHTML = html;
+				return textArea.value;
+			};
+
+			try {
+				// Unescape the HTML-encoded string
+				const unescapedOutput = unescapeHtml(attributes.output);
+
+				// Parse the unescaped string into JSON
+				const output = JSON.parse(unescapedOutput);
+
+				// Assign the parsed values to variables
+				stdout = output.stdout;
+				stderr = output.stderr;
+				result = output.result;
+			} catch (error) {
+				console.error('Error:', error);
+			}
+		}
+	};
+
 	onMount(async () => {
 		console.log('codeblock', lang, code);
 
diff --git a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
index d44e60fa1..cf42bd2ec 100644
--- a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
@@ -23,6 +23,7 @@
 	export let id: string;
 	export let tokens: Token[];
 	export let top = true;
+	export let attributes = {};
 
 	export let save = false;
 	export let onSourceClick: Function = () => {};
@@ -83,6 +84,7 @@
 				{token}
 				lang={token?.lang ?? ''}
 				code={token?.text ?? ''}
+				{attributes}
 				{save}
 				on:code={(e) => {
 					dispatch('code', e.detail);
@@ -197,7 +199,11 @@
 	{:else if token.type === 'details'}
 		<Collapsible title={token.summary} attributes={token?.attributes} className="w-full space-y-1">
 			<div class=" mb-1.5" slot="content">
-				<svelte:self id={`${id}-${tokenIdx}-d`} tokens={marked.lexer(token.text)} />
+				<svelte:self
+					id={`${id}-${tokenIdx}-d`}
+					tokens={marked.lexer(token.text)}
+					attributes={token?.attributes}
+				/>
 			</div>
 		</Collapsible>
 	{:else if token.type === 'html'}
diff --git a/src/lib/components/common/Collapsible.svelte b/src/lib/components/common/Collapsible.svelte
index 0e0c0a4df..4f6b1ea95 100644
--- a/src/lib/components/common/Collapsible.svelte
+++ b/src/lib/components/common/Collapsible.svelte
@@ -80,6 +80,12 @@
 						{:else}
 							{$i18n.t('Thinking...')}
 						{/if}
+					{:else if attributes?.type === 'code_interpreter'}
+						{#if attributes?.done === 'true'}
+							{$i18n.t('Analyzed')}
+						{:else}
+							{$i18n.t('Analyzing...')}
+						{/if}
 					{:else}
 						{title}
 					{/if}
diff --git a/src/lib/workers/pyodide.worker.ts b/src/lib/workers/pyodide.worker.ts
index 7b95683be..1eb426ea3 100644
--- a/src/lib/workers/pyodide.worker.ts
+++ b/src/lib/workers/pyodide.worker.ts
@@ -84,6 +84,10 @@ function processResult(result: any): any {
 			// Handle primitive types directly
 			return result;
 		}
+		if (typeof result === 'bigint') {
+			// Convert BigInt to a string for JSON-safe representation
+			return result.toString();
+		}
 		if (Array.isArray(result)) {
 			// If it's an array, recursively process items
 			return result.map((item) => processResult(item));
@@ -106,7 +110,7 @@ function processResult(result: any): any {
 		return JSON.stringify(result);
 	} catch (err) {
 		// In case something unexpected happens, we return a stringified fallback
-		return `[processResult error]: ${err.toString()}`;
+		return `[processResult error]: ${err.message || err.toString()}`;
 	}
 }