From 5a720a4a31058ed88a58930052128bc6a0ad062e Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Fri, 27 Dec 2024 22:46:50 -0800 Subject: [PATCH] refac --- backend/open_webui/utils/middleware.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index a289b2c59..bb6c56a3b 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -929,6 +929,10 @@ async def process_chat_response( # Handle as a background task async def post_response_handler(response, events): + + assistant_message = get_last_assistant_message(form_data["messages"]) + content = assistant_message if assistant_message else "" + try: for event in events: await event_emitter( @@ -947,9 +951,6 @@ async def process_chat_response( }, ) - assistant_message = get_last_assistant_message(form_data["messages"]) - content = assistant_message if assistant_message else "" - async for line in response.body_iterator: line = line.decode("utf-8") if isinstance(line, bytes) else line data = line @@ -1051,6 +1052,16 @@ async def process_chat_response( print("Task was cancelled!") await event_emitter({"type": "task-cancelled"}) + if not ENABLE_REALTIME_CHAT_SAVE: + # Save message in the database + Chats.upsert_message_to_chat_by_id_and_message_id( + metadata["chat_id"], + metadata["message_id"], + { + "content": content, + }, + ) + if response.background is not None: await response.background()