From 4ef042e966de1013b06490e15ebecd55fcebc5a6 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Fri, 16 Aug 2024 15:33:14 +0200
Subject: [PATCH] refac

---
 .../chat/Messages/MarkdownInlineTokens.svelte |  8 +++
 .../chat/Messages/MarkdownTokens.svelte       |  9 +++
 .../chat/Messages/ResponseMessage.svelte      |  4 +-
 src/lib/utils/index.ts                        |  7 ---
 src/lib/utils/{ => marked}/katex-extension.ts |  0
 src/lib/utils/marked/token-extension.ts       | 57 +++++++++++++++++++
 6 files changed, 77 insertions(+), 8 deletions(-)
 rename src/lib/utils/{ => marked}/katex-extension.ts (100%)
 create mode 100644 src/lib/utils/marked/token-extension.ts

diff --git a/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte b/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
index 8ca6fa37d..c518962ac 100644
--- a/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
@@ -45,6 +45,14 @@
 		{#if token.text}
 			<KatexRenderer content={revertSanitizedResponseContent(token.text)} displayMode={false} />
 		{/if}
+	{:else if token.type === 'iframe'}
+		<iframe
+			src="{WEBUI_BASE_URL}/api/v1/files/{token.fileId}/content"
+			title={token.fileId}
+			width="100%"
+			frameborder="0"
+			onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"
+		></iframe>
 	{:else if token.type === 'text'}
 		{token.raw}
 	{/if}
diff --git a/src/lib/components/chat/Messages/MarkdownTokens.svelte b/src/lib/components/chat/Messages/MarkdownTokens.svelte
index 854bd947f..c183ca861 100644
--- a/src/lib/components/chat/Messages/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownTokens.svelte
@@ -7,6 +7,7 @@
 	import CodeBlock from '$lib/components/chat/Messages/CodeBlock.svelte';
 	import MarkdownInlineTokens from '$lib/components/chat/Messages/MarkdownInlineTokens.svelte';
 	import KatexRenderer from './KatexRenderer.svelte';
+	import { WEBUI_BASE_URL } from '$lib/constants';
 
 	export let id: string;
 	export let tokens: Token[];
@@ -98,6 +99,14 @@
 		{:else}
 			{token.text}
 		{/if}
+	{:else if token.type === 'iframe'}
+		<iframe
+			src="{WEBUI_BASE_URL}/api/v1/files/{token.fileId}/content"
+			title={token.fileId}
+			width="100%"
+			frameborder="0"
+			onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"
+		></iframe>
 	{:else if token.type === 'paragraph'}
 		<p>
 			<MarkdownInlineTokens id={`${id}-${tokenIdx}-p`} tokens={token.tokens ?? []} />
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index c95d0afc5..ce96b7961 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -76,13 +76,15 @@
 
 	import 'katex/dist/katex.min.css';
 
-	import markedKatex from '$lib/utils/katex-extension';
+	import markedKatex from '$lib/utils/marked/katex-extension';
+	import markedToken from '$lib/utils/marked/token-extension';
 
 	const options = {
 		throwOnError: false
 	};
 
 	marked.use(markedKatex(options));
+	marked.use(markedToken({ WEBUI_BASE_URL }));
 
 	$: (async () => {
 		if (message?.content) {
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index f3532773c..c0516deaf 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -27,7 +27,6 @@ export const replaceTokens = (content, char, user) => {
 	const charToken = /{{char}}/gi;
 	const userToken = /{{user}}/gi;
 	const videoIdToken = /{{VIDEO_FILE_ID_([a-f0-9-]+)}}/gi; // Regex to capture the video ID
-	const htmlIdToken = /{{HTML_FILE_ID_([a-f0-9-]+)}}/gi; // Regex to capture the HTML ID
 
 	// Replace {{char}} if char is provided
 	if (char !== undefined && char !== null) {
@@ -45,12 +44,6 @@ export const replaceTokens = (content, char, user) => {
 		return `<video src="${videoUrl}" controls></video>`;
 	});
 
-	// Replace HTML ID tags with corresponding HTML content
-	content = content.replace(htmlIdToken, (match, fileId) => {
-		const htmlUrl = `${WEBUI_BASE_URL}/api/v1/files/${fileId}/content`;
-		return `<iframe src="${htmlUrl}" width="100%" frameborder="0" onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"></iframe>`;
-	});
-
 	return content;
 };
 
diff --git a/src/lib/utils/katex-extension.ts b/src/lib/utils/marked/katex-extension.ts
similarity index 100%
rename from src/lib/utils/katex-extension.ts
rename to src/lib/utils/marked/katex-extension.ts
diff --git a/src/lib/utils/marked/token-extension.ts b/src/lib/utils/marked/token-extension.ts
new file mode 100644
index 000000000..51e28662a
--- /dev/null
+++ b/src/lib/utils/marked/token-extension.ts
@@ -0,0 +1,57 @@
+export default function (options = {}) {
+	return {
+		extensions: [inlineIframeToken(options), blockIframeToken(options)]
+	};
+}
+
+const inlineIframeToken = (options = {}) => {
+	const WEBUI_BASE_URL = options.WEBUI_BASE_URL || '';
+	const htmlIdToken = /{{HTML_FILE_ID_([a-f0-9-]+)}}/gi; // Regex to capture the HTML ID
+	function tokenizer(src) {
+		const match = htmlIdToken.exec(src);
+		if (match) {
+			return {
+				type: 'iframe',
+				raw: match[0],
+				fileId: match[1]
+			};
+		}
+	}
+	function renderer(token) {
+		const htmlUrl = `${WEBUI_BASE_URL}/api/v1/files/${token.fileId}/content`;
+		const iframeElement = `<iframe src="${htmlUrl}" width="100%" frameborder="0" onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"></iframe>`;
+		return iframeElement;
+	}
+	return {
+		name: 'iframe',
+		level: 'inline',
+		tokenizer,
+		renderer
+	};
+};
+
+const blockIframeToken = (options = {}) => {
+	const WEBUI_BASE_URL = options.WEBUI_BASE_URL || '';
+	const htmlIdToken = /{{HTML_FILE_ID_([a-f0-9-]+)}}/gi; // Regex to capture the HTML ID
+	function tokenizer(src) {
+		const match = htmlIdToken.exec(src);
+		if (match) {
+			return {
+				type: 'iframe',
+				raw: match[0],
+				fileId: match[1]
+			};
+		}
+	}
+	function renderer(token) {
+		const htmlUrl = `${WEBUI_BASE_URL}/api/v1/files/${token.fileId}/content`;
+		const iframeElement = `<iframe src="${htmlUrl}" width="100%" frameborder="0" onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"></iframe>`;
+		return iframeElement;
+	}
+	return {
+		name: 'iframe',
+		level: 'block',
+		tokenizer,
+		renderer
+	};
+};