From 75e5a485d200ebfaa6efe5fd6559a999c89c4f93 Mon Sep 17 00:00:00 2001
From: Tim Baek <tim@openwebui.com>
Date: Sun, 1 Feb 2026 10:06:52 +0400
Subject: [PATCH] refac

---
 backend/open_webui/models/chat_messages.py | 44 ++++++++++++++++++++++
 backend/open_webui/utils/middleware.py     | 18 ++++++---
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py
index 202b03c15..4baa7ea20 100644
--- a/backend/open_webui/models/chat_messages.py
+++ b/backend/open_webui/models/chat_messages.py
@@ -269,6 +269,50 @@ class ChatMessageTable:
             results = query.group_by(ChatMessage.model_id).all()
             return {row.model_id: row.count for row in results}
 
+    def get_token_usage_by_model(
+        self,
+        start_date: Optional[int] = None,
+        end_date: Optional[int] = None,
+        db: Optional[Session] = None,
+    ) -> dict[str, dict]:
+        """Aggregate token usage by model. Works with SQLite and PostgreSQL."""
+        with get_db_context(db) as db:
+            query = db.query(
+                ChatMessage.model_id,
+                ChatMessage.usage,
+            ).filter(
+                ChatMessage.role == "assistant",
+                ChatMessage.model_id.isnot(None),
+                ChatMessage.usage.isnot(None),
+            )
+
+            if start_date:
+                query = query.filter(ChatMessage.created_at >= start_date)
+            if end_date:
+                query = query.filter(ChatMessage.created_at <= end_date)
+
+            results = query.all()
+
+            # Aggregate in Python for cross-database compatibility
+            usage_by_model: dict[str, dict] = {}
+            for model_id, usage in results:
+                if model_id not in usage_by_model:
+                    usage_by_model[model_id] = {
+                        "input_tokens": 0,
+                        "output_tokens": 0,
+                        "message_count": 0,
+                    }
+
+                usage_by_model[model_id]["input_tokens"] += usage.get("input_tokens") or 0
+                usage_by_model[model_id]["output_tokens"] += usage.get("output_tokens") or 0
+                usage_by_model[model_id]["message_count"] += 1
+
+            # Add total_tokens
+            for data in usage_by_model.values():
+                data["total_tokens"] = data["input_tokens"] + data["output_tokens"]
+
+            return usage_by_model
+
     def get_message_count_by_user(
         self,
         start_date: Optional[int] = None,
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index c38622663..e4d800ba5 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -3001,6 +3001,7 @@ async def process_chat_response(
                     "content": content,
                 }
             ]
+            usage = None
 
             reasoning_tags_param = metadata.get("params", {}).get("reasoning_tags")
             DETECT_REASONING_TAGS = reasoning_tags_param is not False
@@ -3122,11 +3123,11 @@ async def process_chat_response(
                                 else:
                                     choices = data.get("choices", [])
 
-                                    # 17421 - Normalize usage data to standard format
-                                    usage = data.get("usage", {}) or {}
-                                    usage.update(data.get("timings", {}))  # llama.cpp
-                                    if usage:
-                                        usage = normalize_usage(usage)
+                                    # Normalize usage data to standard format
+                                    raw_usage = data.get("usage", {}) or {}
+                                    raw_usage.update(data.get("timings", {}))  # llama.cpp
+                                    if raw_usage:
+                                        usage = normalize_usage(raw_usage)
                                         await event_emitter(
                                             {
                                                 "type": "chat:completion",
@@ -3924,8 +3925,15 @@ async def process_chat_response(
                         {
                             "content": serialize_output(output),
                             "output": output,
+                            **({"usage": usage} if usage else {}),
                         },
                     )
+                elif usage:
+                    Chats.upsert_message_to_chat_by_id_and_message_id(
+                        metadata["chat_id"],
+                        metadata["message_id"],
+                        {"usage": usage},
+                    )
 
                 # Send a webhook notification if the user is not active
                 if not Users.is_user_active(user.id):