Timothy Jaeryang Baek пре 4 месеци
родитељ
комит
5a720a4a31
1 измењених фајлова са 14 додато и 3 уклоњено
  1. 14 3
      backend/open_webui/utils/middleware.py

+ 14 - 3
backend/open_webui/utils/middleware.py

@@ -929,6 +929,10 @@ async def process_chat_response(
 
 
         # Handle as a background task
         # Handle as a background task
         async def post_response_handler(response, events):
         async def post_response_handler(response, events):
+
+            assistant_message = get_last_assistant_message(form_data["messages"])
+            content = assistant_message if assistant_message else ""
+
             try:
             try:
                 for event in events:
                 for event in events:
                     await event_emitter(
                     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:
                 async for line in response.body_iterator:
                     line = line.decode("utf-8") if isinstance(line, bytes) else line
                     line = line.decode("utf-8") if isinstance(line, bytes) else line
                     data = line
                     data = line
@@ -1051,6 +1052,16 @@ async def process_chat_response(
                 print("Task was cancelled!")
                 print("Task was cancelled!")
                 await event_emitter({"type": "task-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:
             if response.background is not None:
                 await response.background()
                 await response.background()