From 7ff719938a8878446036f7d141f65dd0cdd5d702 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Wed, 12 Feb 2025 23:55:14 -0800
Subject: [PATCH] refac: citations

---
 .../components/chat/Messages/Citations.svelte |  4 +--
 .../chat/Messages/Markdown/Source.svelte      | 25 ++++++++++++-------
 src/lib/utils/index.ts                        |  9 +++----
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/src/lib/components/chat/Messages/Citations.svelte b/src/lib/components/chat/Messages/Citations.svelte
index 041a419bf..b09f51482 100644
--- a/src/lib/components/chat/Messages/Citations.svelte
+++ b/src/lib/components/chat/Messages/Citations.svelte
@@ -100,7 +100,7 @@
 			<div class="flex text-xs font-medium flex-wrap">
 				{#each citations as citation, idx}
 					<button
-						id={`source-${citation.source.name}`}
+						id={`source-${idx}`}
 						class="no-toggle outline-none flex dark:text-gray-300 p-1 bg-white dark:bg-gray-900 rounded-xl max-w-96"
 						on:click={() => {
 							showCitationModal = true;
@@ -179,7 +179,7 @@
 					<div class="flex text-xs font-medium flex-wrap">
 						{#each citations as citation, idx}
 							<button
-								id={`source-${citation.source.name}`}
+								id={`source-${idx}`}
 								class="no-toggle outline-none flex dark:text-gray-300 p-1 bg-gray-50 hover:bg-gray-100 dark:bg-gray-900 dark:hover:bg-gray-850 transition rounded-xl max-w-96"
 								on:click={() => {
 									showCitationModal = true;
diff --git a/src/lib/components/chat/Messages/Markdown/Source.svelte b/src/lib/components/chat/Messages/Markdown/Source.svelte
index 5ff1b8eec..7d64d0ae9 100644
--- a/src/lib/components/chat/Messages/Markdown/Source.svelte
+++ b/src/lib/components/chat/Messages/Markdown/Source.svelte
@@ -2,24 +2,31 @@
 	export let token;
 	export let onClick: Function = () => {};
 
-	let id = '';
-	function extractDataAttribute(input) {
-		// Use a regular expression to extract the value of the `data` attribute
-		const match = input.match(/data="([^"]*)"/);
-		// Check if a match was found and return the first captured group
-		return match ? match[1] : null;
+	let attributes: Record<string, string> = {};
+
+	function extractAttributes(input: string): Record<string, string> {
+		const regex = /(\w+)="([^"]*)"/g;
+		let match;
+		let attrs: Record<string, string> = {};
+
+		// Loop through all matches and populate the attributes object
+		while ((match = regex.exec(input)) !== null) {
+			attrs[match[1]] = match[2];
+		}
+
+		return attrs;
 	}
 
-	$: id = extractDataAttribute(token.text);
+	$: attributes = extractAttributes(token.text);
 </script>
 
 <button
 	class="text-xs font-medium w-fit translate-y-[2px] px-2 py-0.5 dark:bg-white/5 dark:text-white/60 dark:hover:text-white bg-gray-50 text-black/60 hover:text-black transition rounded-lg"
 	on:click={() => {
-		onClick(id);
+		onClick(attributes.data);
 	}}
 >
 	<span class="line-clamp-1">
-		{id}
+		{attributes.title}
 	</span>
 </button>
diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts
index 0a5b72e36..4d1e92f04 100644
--- a/src/lib/utils/index.ts
+++ b/src/lib/utils/index.ts
@@ -55,15 +55,12 @@ export const replaceTokens = (content, sourceIds, char, user) => {
 
 	// Remove sourceIds from the content and replace them with <source_id>...</source_id>
 	if (Array.isArray(sourceIds)) {
-		sourceIds.forEach((sourceId) => {
-			// Escape special characters in the sourceId
-			const escapedSourceId = escapeRegExp(sourceId);
-
+		sourceIds.forEach((sourceId, idx) => {
 			// Create a token based on the exact `[sourceId]` string
-			const sourceToken = `\\[${escapedSourceId}\\]`; // Escape special characters for RegExp
+			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="${sourceId}" />`);
+			content = content.replace(sourceRegex, `<source_id data="${idx}" title="${sourceId}" />`);
 		});
 	}