From 4dea4fdf54e00ebaba8e3178128bf8709453d2a2 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Thu, 2 Apr 2026 08:34:49 -0500
Subject: [PATCH] refac

---
 backend/open_webui/models/chat_messages.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py
index ac75fcf973..b37c04037e 100644
--- a/backend/open_webui/models/chat_messages.py
+++ b/backend/open_webui/models/chat_messages.py
@@ -5,6 +5,7 @@ from typing import Any, Optional
 
 from sqlalchemy.orm import Session
 from open_webui.internal.db import Base, get_db_context
+from open_webui.utils.response import normalize_usage
 
 from pydantic import BaseModel, ConfigDict
 from sqlalchemy import (
@@ -41,6 +42,12 @@ def _normalize_timestamp(timestamp: int) -> float:
     return timestamp
 
 
+def get_usage(data: dict) -> Optional[dict]:
+    """Extract and normalize usage from message data."""
+    usage = data.get('usage') or (data.get('info') or {}).get('usage')
+    return normalize_usage(usage) if usage else None
+
+
 ####################
 # ChatMessage DB Schema
 ####################
@@ -163,11 +170,8 @@ class ChatMessageTable:
                     existing.status_history = data.get('status_history') or data.get('statusHistory')
                 if 'error' in data:
                     existing.error = data.get('error')
-                # 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
+                # Extract and normalize usage
+                usage = get_usage(data)
                 if usage:
                     # Deep-merge: preserve existing keys not present in new data
                     # This prevents background tasks (follow-ups, title, tags)
@@ -179,11 +183,8 @@ class ChatMessageTable:
                 return ChatMessageModel.model_validate(existing)
             else:
                 # Insert new
-                # 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
+                # Extract and normalize usage
+                usage = get_usage(data)
                 message = ChatMessage(
                     id=composite_id,
                     chat_id=chat_id,