diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py index 4baa7ea20..e0f00da56 100644 --- a/backend/open_webui/models/chat_messages.py +++ b/backend/open_webui/models/chat_messages.py @@ -140,9 +140,11 @@ class ChatMessageTable: ) if "error" in data: existing.error = data.get("error") - # Extract usage from info.usage if present - info = data.get("info", {}) - usage = info.get("usage") if info else None + # Extract usage - check direct field first, then info.usage + usage = data.get("usage") + if not usage: + info = data.get("info", {}) + usage = info.get("usage") if info else None if usage: existing.usage = usage existing.updated_at = now @@ -151,9 +153,11 @@ class ChatMessageTable: return ChatMessageModel.model_validate(existing) else: # Insert new - # Extract usage from info.usage if present - info = data.get("info", {}) - usage = info.get("usage") if info else None + # Extract usage - check direct field first, then info.usage + usage = data.get("usage") + if not usage: + info = data.get("info", {}) + usage = info.get("usage") if info else None message = ChatMessage( id=message_id, chat_id=chat_id, diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index e4d800ba5..cf0f14501 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -3042,6 +3042,7 @@ async def process_chat_response( async def stream_body_handler(response, form_data): nonlocal content nonlocal content_blocks + nonlocal usage response_tool_calls = []