From bbda717b69cd26f7a3cb1c87a383e0d5d522a30c Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Sat, 15 Feb 2025 17:24:29 -0800
Subject: [PATCH] refac: citations

---
 backend/open_webui/utils/middleware.py          |  2 +-
 .../chat/Messages/Markdown/Source.svelte        | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index f0cdede3d..52f50b8a9 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -772,7 +772,7 @@ async def process_chat_payload(request, form_data, metadata, user, model):
 
             if "document" in source:
                 for doc_idx, doc_context in enumerate(source["document"]):
-                    context_string += f"<source><source_id>{doc_idx}</source_id><source_context>{doc_context}</source_context></source>\n"
+                    context_string += f"<source><source_id>{source_idx}</source_id><source_context>{doc_context}</source_context></source>\n"
 
         context_string = context_string.strip()
         prompt = get_last_user_message(form_data["messages"])
diff --git a/src/lib/components/chat/Messages/Markdown/Source.svelte b/src/lib/components/chat/Messages/Markdown/Source.svelte
index 7d64d0ae9..4eb1fffb7 100644
--- a/src/lib/components/chat/Messages/Markdown/Source.svelte
+++ b/src/lib/components/chat/Messages/Markdown/Source.svelte
@@ -17,6 +17,21 @@
 		return attrs;
 	}
 
+	// Helper function to return only the domain from a URL
+	function getDomain(url: string): string {
+		const domain = url.replace('http://', '').replace('https://', '').split(/[/?#]/)[0];
+		return domain;
+	}
+
+	// Helper function to check if text is a URL and return the domain
+	function formattedTitle(title: string): string {
+		if (title.startsWith('http')) {
+			return getDomain(title);
+		}
+
+		return title;
+	}
+
 	$: attributes = extractAttributes(token.text);
 </script>
 
@@ -27,6 +42,6 @@
 	}}
 >
 	<span class="line-clamp-1">
-		{attributes.title}
+		{formattedTitle(attributes.title)}
 	</span>
 </button>