refac

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
+                # Extract and normalize usage
-                usage = data.get('usage')
+                usage = get_usage(data)
-                if not usage:
-                    info = data.get('info', {})
-                    usage = info.get('usage') if info else None
                 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
+                # Extract and normalize usage
-                usage = data.get('usage')
+                usage = get_usage(data)
-                if not usage:
-                    info = data.get('info', {})
-                    usage = info.get('usage') if info else None
                 message = ChatMessage(
                     id=composite_id,
                     chat_id=chat_id,