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

---
 .../chat/Messages/MarkdownInlineTokens.svelte |  2 +
 .../chat/Messages/MarkdownTokens.svelte       |  4 +-
 .../chat/Messages/ResponseMessage.svelte      |  3 -
 src/lib/utils/index.ts                        |  7 +++
 src/lib/utils/marked/token-extension.ts       | 57 -------------------
 5 files changed, 12 insertions(+), 61 deletions(-)
 delete 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 9fca021f9..3ee87114c 100644
--- a/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownInlineTokens.svelte
@@ -18,6 +18,8 @@
 		{@const html = DOMPurify.sanitize(token.text)}
 		{#if html && html.includes('<video')}
 			{@html html}
+		{:else if token.text.includes(`<iframe src="${WEBUI_BASE_URL}/api/v1/files/`)}
+			{@html `${token.text}`}
 		{:else}
 			{token.text}
 		{/if}
diff --git a/src/lib/components/chat/Messages/MarkdownTokens.svelte b/src/lib/components/chat/Messages/MarkdownTokens.svelte
index 5caa2b96e..e51dcc2ee 100644
--- a/src/lib/components/chat/Messages/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/MarkdownTokens.svelte
@@ -1,7 +1,7 @@
 <script lang="ts">
 	import DOMPurify from 'dompurify';
 	import { onMount } from 'svelte';
-	import type { Token } from 'marked';
+	import { marked, type Token } from 'marked';
 	import { revertSanitizedResponseContent, unescapeHtml } from '$lib/utils';
 
 	import CodeBlock from '$lib/components/chat/Messages/CodeBlock.svelte';
@@ -96,6 +96,8 @@
 		{@const html = DOMPurify.sanitize(token.text)}
 		{#if html && html.includes('<video')}
 			{@html html}
+		{:else if token.text.includes(`<iframe src="${WEBUI_BASE_URL}/api/v1/files/`)}
+			{@html `${token.text}`}
 		{:else}
 			{token.text}
 		{/if}
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index a7edd3fc9..515caa2e4 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -75,16 +75,13 @@
 	let tokens;
 
 	import 'katex/dist/katex.min.css';
-
 	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 c0516deaf..f3532773c 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -27,6 +27,7 @@ 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) {
@@ -44,6 +45,12 @@ 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/marked/token-extension.ts b/src/lib/utils/marked/token-extension.ts
deleted file mode 100644
index 51e28662a..000000000
--- a/src/lib/utils/marked/token-extension.ts
+++ /dev/null
@@ -1,57 +0,0 @@
-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
-	};
-};