diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index a5f848b62..19e949e88 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -311,6 +311,11 @@ RESET_CONFIG_ON_START = ( os.environ.get("RESET_CONFIG_ON_START", "False").lower() == "true" ) + +ENABLE_REALTIME_CHAT_SAVE = ( + os.environ.get("ENABLE_REALTIME_CHAT_SAVE", "True").lower() == "true" +) + #################################### # REDIS #################################### diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 7d79e1d0b..bb6c56a3b 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -65,6 +65,7 @@ from open_webui.env import ( SRC_LOG_LEVELS, GLOBAL_LOG_LEVEL, BYPASS_MODEL_ACCESS_CONTROL, + ENABLE_REALTIME_CHAT_SAVE, ) from open_webui.constants import TASKS @@ -928,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( @@ -946,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 @@ -977,7 +979,6 @@ async def process_chat_response( ) else: - value = ( data.get("choices", [])[0] .get("delta", {}) @@ -987,6 +988,28 @@ async def process_chat_response( if value: content = f"{content}{value}" + if 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, + }, + ) + else: + data = { + "content": content, + } + + except Exception as e: + done = "data: [DONE]" in line + title = Chats.get_chat_title_by_id(metadata["chat_id"]) + + if done: + data = {"done": True, "content": content, "title": title} + + if not ENABLE_REALTIME_CHAT_SAVE: # Save message in the database Chats.upsert_message_to_chat_by_id_and_message_id( metadata["chat_id"], @@ -996,13 +1019,6 @@ async def process_chat_response( }, ) - except Exception as e: - done = "data: [DONE]" in line - title = Chats.get_chat_title_by_id(metadata["chat_id"]) - - if done: - data = {"done": True, "content": content, "title": title} - # Send a webhook notification if the user is not active if ( get_user_id_from_session_pool(metadata["session_id"]) @@ -1036,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() diff --git a/src/lib/components/channel/Messages.svelte b/src/lib/components/channel/Messages.svelte index 3614263a1..deb2d93fb 100644 --- a/src/lib/components/channel/Messages.svelte +++ b/src/lib/components/channel/Messages.svelte @@ -66,7 +66,7 @@ {($settings?.widescreenMode ?? null) ? 'max-w-full' : 'max-w-5xl'} mx-auto" > {#if channel} -