From 7d42a79177aef9fc12ede4f0bbeccfe3ffe53a3f Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 5 Aug 2024 17:12:41 +0200
Subject: [PATCH] refac: token rendering

---
 .../chat/Messages/HTMLRenderer.svelte         | 30 ++++++
 .../chat/Messages/TokenRenderer.svelte        | 97 +++++++++++--------
 src/lib/components/common/Image.svelte        |  4 +-
 3 files changed, 90 insertions(+), 41 deletions(-)
 create mode 100644 src/lib/components/chat/Messages/HTMLRenderer.svelte

diff --git a/src/lib/components/chat/Messages/HTMLRenderer.svelte b/src/lib/components/chat/Messages/HTMLRenderer.svelte
new file mode 100644
index 000000000..53fad8084
--- /dev/null
+++ b/src/lib/components/chat/Messages/HTMLRenderer.svelte
@@ -0,0 +1,30 @@
+<script lang="ts">
+	import Image from '$lib/components/common/Image.svelte';
+	import CodeBlock from './CodeBlock.svelte';
+
+	/* The html content of the tag */
+	export let html; //: string;
+	let parsedHTML = [html];
+
+	export let images;
+	export let codes;
+
+	//  all images are in {{IMAGE_0}}, {{IMAGE_1}}....  format
+	// all codes are in {{CODE_0}}, {{CODE_1}}....  format
+
+	const rules = [];
+	rules.forEach((rule) => {
+		parsedHTML = parsedHTML.map((substr) => substr.split(rule.regex)).flat();
+	});
+</script>
+
+{#each parsedHTML as part}
+	{@const match = rules.find((rule) => rule.regex.test(part))}
+	{#if match}
+		<svelte:component this={match.component} {...match.props}>
+			{@html part}
+		</svelte:component>
+	{:else}
+		{@html part}
+	{/if}
+{/each}
diff --git a/src/lib/components/chat/Messages/TokenRenderer.svelte b/src/lib/components/chat/Messages/TokenRenderer.svelte
index ecef0fcb4..13fd56bea 100644
--- a/src/lib/components/chat/Messages/TokenRenderer.svelte
+++ b/src/lib/components/chat/Messages/TokenRenderer.svelte
@@ -10,6 +10,9 @@
 	export let tokenIdx = 0;
 	export let id;
 
+	let element;
+	let content;
+
 	const renderer = new marked.Renderer();
 
 	// For code blocks with simple backticks
@@ -17,30 +20,19 @@
 		return `<code>${code.replaceAll('&amp;', '&')}</code>`;
 	};
 
-	// renderer.code = (code, lang) => {
-	// 	const element = document.createElement('div');
-	// 	new CodeBlock({
-	// 		target: element,
-	// 		props: {
-	// 			id: `${id}-${tokenIdx}`,
-	// 			lang: lang ?? '',
-	// 			code: revertSanitizedResponseContent(code ?? '')
-	// 		}
-	// 	});
-	// 	return element.innerHTML;
-	// };
+	let codes = [];
+	renderer.code = (code, lang) => {
+		codes.push({ code, lang, id: codes.length });
+		codes = codes;
+		return `{{@CODE ${codes.length - 1}}}`;
+	};
 
-	// renderer.image = (href, title, text) => {
-	// 	const element = document.createElement('pre');
-	// 	new Image({
-	// 		target: element,
-	// 		props: {
-	// 			src: href,
-	// 			alt: text
-	// 		}
-	// 	});
-	// 	return element.innerHTML;
-	// };
+	let images = [];
+	renderer.image = (href, title, text) => {
+		images.push({ href, title, text });
+		images = images;
+		return `{{@IMAGE ${images.length - 1}}}`;
+	};
 
 	// Open all links in a new tab/window (from https://github.com/markedjs/marked/issues/655#issuecomment-383226346)
 	const origLinkRenderer = renderer.link;
@@ -53,23 +45,50 @@
 		// eslint-disable-next-line @typescript-eslint/no-explicit-any
 		extensions: any;
 	};
+
+	$: if (token) {
+		images = [];
+		codes = [];
+		content = marked
+			.parse(token.raw, {
+				...defaults,
+				gfm: true,
+				breaks: true,
+				renderer
+			})
+			.split(/({{@IMAGE [^}]+}}|{{@CODE [^}]+}})/g);
+	}
 </script>
 
-{#if token.type === 'code'}
-	{#if token.lang === 'mermaid'}
-		<pre class="mermaid">{revertSanitizedResponseContent(token.text)}</pre>
+<div bind:this={element}>
+	{#if token.type === 'code'}
+		{#if token.lang === 'mermaid'}
+			<pre class="mermaid">{revertSanitizedResponseContent(token.text)}</pre>
+		{:else}
+			<CodeBlock
+				id={`${id}-${tokenIdx}`}
+				lang={token?.lang ?? ''}
+				code={revertSanitizedResponseContent(token?.text ?? '')}
+			/>
+		{/if}
+	{:else if token.type === 'image'}
+		<Image src={token.href} alt={token.text} />
 	{:else}
-		<CodeBlock
-			id={`${id}-${tokenIdx}`}
-			lang={token?.lang ?? ''}
-			code={revertSanitizedResponseContent(token?.text ?? '')}
-		/>
+		{#each content as part}
+			{@html part.startsWith('{{@IMAGE ') || part.startsWith('{{@CODE ') ? '' : part}
+
+			{#if images.length > 0 && part.startsWith('{{@IMAGE ')}
+				{@const img = images[parseInt(part.match(/{{@IMAGE (\d+)}}/)[1])]}
+
+				<div class="mt-6">
+					<Image src={img.href} text={img.text} />
+				</div>
+			{:else if codes.length > 0 && part.startsWith('{{@CODE ')}
+				{@const _code = codes[parseInt(part.match(/{{@CODE (\d+)}}/)[1])]}
+				<div class="my-10 -mb-6">
+					<CodeBlock id={`${id}-${tokenIdx}-${_code.id}`} lang={_code.lang} code={_code.code} />
+				</div>
+			{/if}
+		{/each}
 	{/if}
-{:else}
-	{@html marked.parse(token.raw, {
-		...defaults,
-		gfm: true,
-		breaks: true,
-		renderer
-	})}
-{/if}
+</div>
diff --git a/src/lib/components/common/Image.svelte b/src/lib/components/common/Image.svelte
index 435ad64ee..425b9bc4f 100644
--- a/src/lib/components/common/Image.svelte
+++ b/src/lib/components/common/Image.svelte
@@ -12,14 +12,14 @@
 	let showImagePreview = false;
 </script>
 
-<div class="w-full py-3">
+<div>
 	<button
 		on:click={() => {
 			showImagePreview = true;
 		}}
+		class="w-full"
 	>
 		<img src={_src} {alt} class=" max-h-96 rounded-lg" draggable="false" data-cy="image" />
 	</button>
 </div>
-
 <ImagePreview bind:show={showImagePreview} src={_src} {alt} />