From d97a4d687ebc065ae4ec2a45b24c775e052ab38c Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 1 Jul 2024 21:41:44 -0700
Subject: [PATCH] refac

---
 backend/main.py                     | 17 +++++++++++++++++
 src/lib/components/chat/Chat.svelte | 10 ++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/backend/main.py b/backend/main.py
index 7731df2f8..0e3986f21 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -1123,6 +1123,17 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
             else:
                 pass
 
+    async def __event_emitter__(data):
+        await sio.emit(
+            "chat-events",
+            {
+                "chat_id": data["chat_id"],
+                "message_id": data["id"],
+                "data": data,
+            },
+            to=data["session_id"],
+        )
+
     def get_priority(function_id):
         function = Functions.get_function_by_id(function_id)
         if function is not None and hasattr(function, "valves"):
@@ -1204,6 +1215,12 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
                             "__model__": model,
                         }
 
+                    if "__event_emitter__" in sig.parameters:
+                        params = {
+                            **params,
+                            "__event_emitter__": __event_emitter__,
+                        }
+
                     if inspect.iscoroutinefunction(outlet):
                         data = await outlet(**params)
                     else:
diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte
index f077833fb..3d03246b7 100644
--- a/src/lib/components/chat/Chat.svelte
+++ b/src/lib/components/chat/Chat.svelte
@@ -327,7 +327,7 @@
 		}
 	};
 
-	const chatCompletedHandler = async (modelId, messages) => {
+	const chatCompletedHandler = async (modelId, responseMessageId, messages) => {
 		await mermaid.run({
 			querySelector: '.mermaid'
 		});
@@ -341,7 +341,9 @@
 				info: m.info ? m.info : undefined,
 				timestamp: m.timestamp
 			})),
-			chat_id: $chatId
+			chat_id: $chatId,
+			session_id: $socket?.id,
+			id: responseMessageId
 		}).catch((error) => {
 			toast.error(error);
 			messages.at(-1).error = { content: error };
@@ -731,7 +733,7 @@
 						controller.abort('User: Stop Response');
 					} else {
 						const messages = createMessagesList(responseMessageId);
-						await chatCompletedHandler(model.id, messages);
+						await chatCompletedHandler(model.id, responseMessageId, messages);
 					}
 
 					_response = responseMessage.content;
@@ -1041,7 +1043,7 @@
 						} else {
 							const messages = createMessagesList(responseMessageId);
 
-							await chatCompletedHandler(model.id, messages);
+							await chatCompletedHandler(model.id, responseMessageId, messages);
 						}
 
 						_response = responseMessage.content;