From 4c6567c46fcaf3ac4c6ea33c143820c8d8389e55 Mon Sep 17 00:00:00 2001
From: Jun Siang Cheah <me@jscheah.me>
Date: Mon, 6 May 2024 22:06:37 +0800
Subject: [PATCH] feat: group citations by source

---
 .../chat/Messages/CitationsModal.svelte       |  2 -
 .../chat/Messages/ResponseMessage.svelte      | 37 ++++++++++++++++---
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/lib/components/chat/Messages/CitationsModal.svelte b/src/lib/components/chat/Messages/CitationsModal.svelte
index 6ed4125ca..79a9a9720 100644
--- a/src/lib/components/chat/Messages/CitationsModal.svelte
+++ b/src/lib/components/chat/Messages/CitationsModal.svelte
@@ -49,7 +49,6 @@
 
 		<div class="flex flex-col w-full px-5 py-4 dark:text-gray-200 overflow-y-scroll max-h-[22rem]">
 			{#each mergedDocuments as document}
-				<!-- Source from document.metadata.source -->
 				<div class="flex flex-col w-full">
 					<div class="text-lg font-medium dark:text-gray-300">
 						{$i18n.t('Source')}
@@ -58,7 +57,6 @@
 						{document.metadata?.source ?? $i18n.t('No source available')}
 					</div>
 				</div>
-				<!-- Content from document.document.content -->
 				<div class="flex flex-col w-full">
 					<div class="text-lg font-medium dark:text-gray-300">
 						{$i18n.t('Content')}
diff --git a/src/lib/components/chat/Messages/ResponseMessage.svelte b/src/lib/components/chat/Messages/ResponseMessage.svelte
index d7da104fe..de59239e1 100644
--- a/src/lib/components/chat/Messages/ResponseMessage.svelte
+++ b/src/lib/components/chat/Messages/ResponseMessage.svelte
@@ -67,6 +67,8 @@
 	let showRateComment = false;
 
 	let showCitations = {};
+	// Backend returns a list of citations per collection, we flatten it to citations per source
+	let flattenedCitations = {};
 
 	$: tokens = marked.lexer(sanitizeResponseContent(message.content));
 
@@ -133,6 +135,28 @@
 				allowHTML: true
 			});
 		}
+
+		if (message.citations) {
+			for (const citation of message.citations) {
+				const zipped = (citation?.document ?? []).map(function (document, index) {
+					return [document, citation.metadata?.[index]];
+				});
+				for (const [document, metadata] of zipped) {
+					const source = metadata?.source ?? 'N/A';
+					if (source in flattenedCitations) {
+						flattenedCitations[source].document.push(document);
+						flattenedCitations[source].metadata.push(metadata);
+					} else {
+						flattenedCitations[source] = {
+							document: [document],
+							metadata: [metadata]
+						};
+					}
+				}
+			}
+			console.log(flattenedCitations);
+			console.log(Object.keys(flattenedCitations));
+		}
 	};
 
 	const renderLatex = () => {
@@ -363,16 +387,19 @@
 						{/each}
 					</div>
 				{/if}
-				{#if message.citations}
+				{#if flattenedCitations}
 					<div class="my-2.5 w-full flex overflow-x-auto gap-2 flex-wrap">
-						{#each message.citations as citation}
+						{#each [...Object.keys(flattenedCitations)] as source}
 							<div>
-								<CitationsModal bind:show={showCitations[citation]} {citation} />
+								<CitationsModal
+									bind:show={showCitations[source]}
+									citation={flattenedCitations[source]}
+								/>
 								<button
 									class="h-16 w-[15rem] flex items-center space-x-3 px-2.5 dark:bg-gray-600 rounded-xl border border-gray-200 dark:border-none text-left"
 									type="button"
 									on:click={() => {
-										showCitations[citation] = !showCitations[citation];
+										showCitations[source] = !showCitations[source];
 									}}
 								>
 									<div class="p-2.5 bg-red-400 text-white rounded-lg">
@@ -395,7 +422,7 @@
 
 									<div class="flex flex-col justify-center -space-y-0.5">
 										<div class=" dark:text-gray-100 text-sm font-medium line-clamp-1">
-											{citation.metadata?.[0]?.source ?? 'N/A'}
+											{source}
 										</div>
 
 										<div class=" text-gray-500 text-sm">{$i18n.t('Document')}</div>