From d844fc7edb6314bf2422ea1b8bbd1c360f694f2e Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Mon, 3 Mar 2025 19:48:00 -0800
Subject: [PATCH] refac: inline citation handling

---
 .../Messages/Markdown/MarkdownTokens.svelte   |  3 +
 src/lib/utils/index.ts                        | 75 ++++++++++---------
 2 files changed, 44 insertions(+), 34 deletions(-)

diff --git a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
index baecd8cd6..8de359a26 100644
--- a/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
+++ b/src/lib/components/chat/Messages/Markdown/MarkdownTokens.svelte
@@ -17,6 +17,7 @@
 	import Collapsible from '$lib/components/common/Collapsible.svelte';
 	import Tooltip from '$lib/components/common/Tooltip.svelte';
 	import ArrowDownTray from '$lib/components/icons/ArrowDownTray.svelte';
+	import Source from './Source.svelte';
 
 	const dispatch = createEventDispatcher();
 
@@ -261,6 +262,8 @@
 			{@html html}
 		{:else if token.text.includes(`<iframe src="${WEBUI_BASE_URL}/api/v1/files/`)}
 			{@html `${token.text}`}
+		{:else if token.text.includes(`<source_id`)}
+			<Source {id} {token} onClick={onSourceClick} />
 		{:else}
 			{token.text}
 		{/if}
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index ef612c4d1..3e905fe53 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -26,43 +26,50 @@ function escapeRegExp(string: string): string {
 }
 
 export const replaceTokens = (content, sourceIds, 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
+	const tokens = [
+		{ regex: /{{char}}/gi, replacement: char },
+		{ regex: /{{user}}/gi, replacement: user },
+		{
+			regex: /{{VIDEO_FILE_ID_([a-f0-9-]+)}}/gi,
+			replacement: (_, fileId) =>
+				`<video src="${WEBUI_BASE_URL}/api/v1/files/${fileId}/content" controls></video>`
+		},
+		{
+			regex: /{{HTML_FILE_ID_([a-f0-9-]+)}}/gi,
+			replacement: (_, fileId) =>
+				`<iframe src="${WEBUI_BASE_URL}/api/v1/files/${fileId}/content/html" width="100%" frameborder="0" onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"></iframe>`
+		}
+	];
 
-	// Replace {{char}} if char is provided
-	if (char !== undefined && char !== null) {
-		content = content.replace(charToken, char);
-	}
+	// Replace tokens outside code blocks only
+	const processOutsideCodeBlocks = (text, replacementFn) => {
+		return text
+			.split(/(```[\s\S]*?```|`[\s\S]*?`)/)
+			.map((segment) => {
+				return segment.startsWith('```') || segment.startsWith('`')
+					? segment
+					: replacementFn(segment);
+			})
+			.join('');
+	};
 
-	// Replace {{user}} if user is provided
-	if (user !== undefined && user !== null) {
-		content = content.replace(userToken, user);
-	}
-
-	// Replace video ID tags with corresponding <video> elements
-	content = content.replace(videoIdToken, (match, fileId) => {
-		const videoUrl = `${WEBUI_BASE_URL}/api/v1/files/${fileId}/content`;
-		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/html`;
-		return `<iframe src="${htmlUrl}" width="100%" frameborder="0" onload="this.style.height=(this.contentWindow.document.body.scrollHeight+20)+'px';"></iframe>`;
-	});
-
-	// Remove sourceIds from the content and replace them with <source_id>...</source_id>
-	if (Array.isArray(sourceIds)) {
-		sourceIds.forEach((sourceId, idx) => {
-			// Create a token based on the exact `[sourceId]` string
-			const sourceToken = `\\[${idx}\\]`; // Escape special characters for RegExp
-			const sourceRegex = new RegExp(sourceToken, 'g'); // Match all occurrences of [sourceId]
-
-			content = content.replace(sourceRegex, `<source_id data="${idx}" title="${sourceId}" />`);
+	// Apply replacements
+	content = processOutsideCodeBlocks(content, (segment) => {
+		tokens.forEach(({ regex, replacement }) => {
+			if (replacement !== undefined && replacement !== null) {
+				segment = segment.replace(regex, replacement);
+			}
 		});
-	}
+
+		if (Array.isArray(sourceIds)) {
+			sourceIds.forEach((sourceId, idx) => {
+				const regex = new RegExp(`\\[${idx}\\]`, 'g');
+				segment = segment.replace(regex, `<source_id data="${idx}" title="${sourceId}" />`);
+			});
+		}
+
+		return segment;
+	});
 
 	return content;
 };