UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-01-31 15:01:00 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek 2024-12-28 19:31:03 -08:00
parent 091e9e2f40
commit 36d927e1ba
3 changed files with 41 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -212,6 +212,15 @@ class ChatTable:
return chat.chat.get("history", {}).get("messages", {}) or {}
def get_message_by_id_and_message_id(
self, id: str, message_id: str
) -> Optional[dict]:
chat = self.get_chat_by_id(id)
if chat is None:
return None
return chat.chat.get("history", {}).get("messages", {}).get(message_id, {})
def upsert_message_to_chat_by_id_and_message_id(
self, id: str, message_id: str, message: dict
) -> Optional[ChatModel]:

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

@ -292,6 +292,34 @@ def get_event_emitter(request_info):
event_data.get("data", {}),
)
if "type" in event_data and event_data["type"] == "message":
message = Chats.get_message_by_id_and_message_id(
request_info["chat_id"],
request_info["message_id"],
)
content = message.get("content", "")
content += event_data.get("data", {}).get("content", "")
Chats.upsert_message_to_chat_by_id_and_message_id(
request_info["chat_id"],
request_info["message_id"],
{
"content": content,
},
)
if "type" in event_data and event_data["type"] == "replace":
content = event_data.get("data", {}).get("content", "")
Chats.upsert_message_to_chat_by_id_and_message_id(
request_info["chat_id"],
request_info["message_id"],
{
"content": content,
},
)
return __event_emitter__

7
backend/open_webui/utils/middleware.py
View File

@ -929,9 +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 ""
message = Chats.get_message_by_id_and_message_id(
metadata["chat_id"], metadata["message_id"]
)
content = message.get("content", "") if message else ""
try:
for event in events: