refac

Author: tjbck

backend/open_webui/utils/middleware.py

@@ -3027,6 +3027,7 @@ async def non_streaming_chat_response_handler(response, ctx):
                         metadata['chat_id'],
                         metadata['message_id'],
                         {
+                            'done': True,
                             'role': 'assistant',
                             'content': content,
                             'output': response_output,
@@ -4376,6 +4377,7 @@ def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
                         metadata['chat_id'],
                         metadata['message_id'],
                         {
+                            'done': True,
                             'content': serialize_output(output),
                             'output': output,
                             **({'usage': usage} if usage else {}),
@@ -4385,7 +4387,13 @@ def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
                     Chats.upsert_message_to_chat_by_id_and_message_id(
                         metadata['chat_id'],
                         metadata['message_id'],
-                        {'usage': usage},
+                        {'done': True, 'usage': usage},
+                    )
+                else:
+                    Chats.upsert_message_to_chat_by_id_and_message_id(
+                        metadata['chat_id'],
+                        metadata['message_id'],
+                        {'done': True},
                     )
 
                 # Send a webhook notification if the user is not active
@@ -4422,10 +4430,17 @@ def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
                         metadata['chat_id'],
                         metadata['message_id'],
                         {
+                            'done': True,
                             'content': serialize_output(output),
                             'output': output,
                         },
                     )
+                else:
+                    Chats.upsert_message_to_chat_by_id_and_message_id(
+                        metadata['chat_id'],
+                        metadata['message_id'],
+                        {'done': True},
+                    )
 
             if response.background is not None:
                 await response.background()

src/lib/components/chat/Chat.svelte

@@ -1239,7 +1239,7 @@
 
 				if (history.currentId) {
 					for (const message of Object.values(history.messages)) {
-						if (message && message.role === 'assistant') {
+						if (message && message.role === 'assistant' && message.done !== false) {
 							message.done = true;
 						}
 					}