From d9573befffd1acdcb87c700a9f566cecdc1f2e05 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Thu, 19 Dec 2024 15:14:09 -0800
Subject: [PATCH] feat: chat completion notification

---
 backend/open_webui/models/chats.py          |  7 +++
 backend/open_webui/utils/middleware.py      | 57 +++++++++++++--------
 src/lib/components/NotificationToast.svelte | 34 ++++++++++++
 src/lib/components/chat/Chat.svelte         | 19 ++++++-
 src/routes/+layout.svelte                   |  2 +-
 5 files changed, 96 insertions(+), 23 deletions(-)
 create mode 100644 src/lib/components/NotificationToast.svelte

diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py
index 75e5114b8..8b2766a1f 100644
--- a/backend/open_webui/models/chats.py
+++ b/backend/open_webui/models/chats.py
@@ -198,6 +198,13 @@ class ChatTable:
             self.add_chat_tag_by_id_and_user_id_and_tag_name(id, user.id, tag_name)
         return self.get_chat_by_id(id)
 
+    def get_chat_title_by_id(self, id: str) -> Optional[str]:
+        chat = self.get_chat_by_id(id)
+        if chat is None:
+            return None
+
+        return chat.chat.get("title", "New Chat")
+
     def get_messages_by_chat_id(self, id: str) -> Optional[dict]:
         chat = self.get_chat_by_id(id)
         if chat is None:
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index f187bd866..6d46b738d 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -582,9 +582,10 @@ async def process_chat_response(request, response, user, events, metadata, tasks
 
                     except Exception as e:
                         done = "data: [DONE]" in line
+                        title = Chats.get_chat_title_by_id(metadata["chat_id"])
 
                         if done:
-                            data = {"done": True}
+                            data = {"done": True, "content": content, "title": title}
                         else:
                             continue
 
@@ -602,27 +603,41 @@ async def process_chat_response(request, response, user, events, metadata, tasks
                     messages = get_message_list(message_map, message.get("id"))
 
                     if tasks:
-                        if (
-                            TASKS.TITLE_GENERATION in tasks
-                            and tasks[TASKS.TITLE_GENERATION]
-                        ):
-                            res = await generate_title(
-                                request,
-                                {
-                                    "model": message["model"],
-                                    "messages": messages,
-                                    "chat_id": metadata["chat_id"],
-                                },
-                                user,
-                            )
-
-                            if res and isinstance(res, dict):
-                                title = (
-                                    res.get("choices", [])[0]
-                                    .get("message", {})
-                                    .get("content", message.get("content", "New Chat"))
+                        if TASKS.TITLE_GENERATION in tasks:
+                            if tasks[TASKS.TITLE_GENERATION]:
+                                res = await generate_title(
+                                    request,
+                                    {
+                                        "model": message["model"],
+                                        "messages": messages,
+                                        "chat_id": metadata["chat_id"],
+                                    },
+                                    user,
                                 )
 
+                                if res and isinstance(res, dict):
+                                    title = (
+                                        res.get("choices", [])[0]
+                                        .get("message", {})
+                                        .get(
+                                            "content",
+                                            message.get("content", "New Chat"),
+                                        )
+                                    )
+
+                                    Chats.update_chat_title_by_id(
+                                        metadata["chat_id"], title
+                                    )
+
+                                    await event_emitter(
+                                        {
+                                            "type": "chat:title",
+                                            "data": title,
+                                        }
+                                    )
+                            elif len(messages) == 2:
+                                title = messages[0].get("content", "New Chat")
+
                                 Chats.update_chat_title_by_id(
                                     metadata["chat_id"], title
                                 )
@@ -630,7 +645,7 @@ async def process_chat_response(request, response, user, events, metadata, tasks
                                 await event_emitter(
                                     {
                                         "type": "chat:title",
-                                        "data": title,
+                                        "data": message.get("content", "New Chat"),
                                     }
                                 )
 
diff --git a/src/lib/components/NotificationToast.svelte b/src/lib/components/NotificationToast.svelte
new file mode 100644
index 000000000..0cc951129
--- /dev/null
+++ b/src/lib/components/NotificationToast.svelte
@@ -0,0 +1,34 @@
+<script lang="ts">
+	import DOMPurify from 'dompurify';
+
+	import { marked } from 'marked';
+	import { createEventDispatcher } from 'svelte';
+
+	const dispatch = createEventDispatcher();
+
+	export let onClick: Function = () => {};
+	export let title: string = 'HI';
+	export let content: string;
+</script>
+
+<button
+	class="flex gap-3 text-left w-[var(--width)] dark:bg-gray-850 dark:text-white bg-white text-black border border-gray-50 dark:border-gray-800 rounded-xl px-3 py-4"
+	on:click={() => {
+		onClick();
+		dispatch('closeToast');
+	}}
+>
+	<div class="flex-shrink-0 self-top -translate-y-1">
+		<img src={'/static/favicon.png'} alt="favicon" class="size-8 rounded-full" />
+	</div>
+
+	<div>
+		{#if title}
+			<div class=" text-[13px] font-medium mb-0.5 line-clamp-1">{title}</div>
+		{/if}
+
+		<div class=" line-clamp-2 text-xs self-center dark:text-gray-300 font-normal">
+			{@html DOMPurify.sanitize(marked(content))}
+		</div>
+	</div>
+</button>
diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte
index be2bd3566..8d64ecbad 100644
--- a/src/lib/components/chat/Chat.svelte
+++ b/src/lib/components/chat/Chat.svelte
@@ -79,6 +79,7 @@
 	import EventConfirmDialog from '../common/ConfirmDialog.svelte';
 	import Placeholder from './Placeholder.svelte';
 	import { getTools } from '$lib/apis/tools';
+	import NotificationToast from '../NotificationToast.svelte';
 
 	export let chatIdProp = '';
 
@@ -308,7 +309,23 @@
 			const type = event?.data?.type ?? null;
 			const data = event?.data?.data ?? null;
 
-			if (type === 'chat:title') {
+			if (type === 'chat:completion') {
+				const { done, content, title } = data;
+
+				if (done) {
+					toast.custom(NotificationToast, {
+						componentProps: {
+							onClick: () => {
+								goto(`/c/${event.chat_id}`);
+							},
+							content: content,
+							title: title
+						},
+						duration: 15000,
+						unstyled: true
+					});
+				}
+			} else if (type === 'chat:title') {
 				currentChatPage.set(1);
 				await chats.set(await getChatList(localStorage.token, $currentChatPage));
 			} else if (type === 'chat:tags') {
diff --git a/src/routes/+layout.svelte b/src/routes/+layout.svelte
index d92e8c2ab..1066cb675 100644
--- a/src/routes/+layout.svelte
+++ b/src/routes/+layout.svelte
@@ -219,5 +219,5 @@
 				: 'light'
 			: 'light'}
 	richColors
-	position="top-center"
+	position="top-right"
 />