UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-12-28 06:42:47 +00:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

enh: save status from the backend

Browse Source
This commit is contained in:
Timothy Jaeryang Baek 2024-12-24 18:03:14 -07:00
parent 95da0734b6
commit 55da6224b8
2 changed files with 27 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
backend/open_webui/models/chats.py
View File

@ -235,6 +235,24 @@ class ChatTable:
chat["history"] = history
return self.update_chat_by_id(id, chat)
def add_message_status_to_chat_by_id_and_message_id(
self, id: str, message_id: str, status: dict
) -> Optional[ChatModel]:
chat = self.get_chat_by_id(id)
if chat is None:
return None
chat = chat.chat
history = chat.get("history", {})
if message_id in history.get("messages", {}):
status_history = history["messages"][message_id].get("statusHistory", [])
status_history.append(status)
history["messages"][message_id]["statusHistory"] = status_history
chat["history"] = history
return self.update_chat_by_id(id, chat)
def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
with get_db() as db:
# Get the existing chat to share

9
backend/open_webui/socket/main.py
View File

@ -6,6 +6,8 @@ import time
from open_webui.models.users import Users
from open_webui.models.channels import Channels
from open_webui.models.chats import Chats
from open_webui.env import (
ENABLE_WEBSOCKET_SUPPORT,
WEBSOCKET_MANAGER,
@ -259,6 +261,13 @@ def get_event_emitter(request_info):
to=session_id,
)
if "type" in event_data and event_data["type"] == "status":
Chats.add_message_status_to_chat_by_id_and_message_id(
request_info["chat_id"],
request_info["message_id"],
event_data.get("data", {}),
)
return __event_emitter__