refac

Author: tjbck

backend/open_webui/models/chat_messages.py

@@ -6,7 +6,7 @@
 from sqlalchemy import select, delete, func, cast, Integer, distinct
 from sqlalchemy.ext.asyncio import AsyncSession
 from open_webui.internal.db import Base, get_async_db_context
-from open_webui.utils.response import normalize_usage
+from open_webui.utils.response import merge_usage, normalize_usage
 from pydantic import BaseModel, ConfigDict
 from sqlalchemy import (
     JSON,
@@ -203,10 +203,8 @@ async def upsert_message(
                 # 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)
-                    # from accidentally clearing the primary response's token counts
-                    existing.usage = {**(existing.usage or {}), **usage}
+                    existing_usage = normalize_usage(existing.usage or {}) if existing.usage else {}
+                    existing.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage)
                 existing.updated_at = now
                 await db.commit()
                 await db.refresh(existing)

backend/open_webui/utils/middleware.py

@@ -112,7 +112,7 @@
 )
 from open_webui.utils.payload import apply_system_prompt_to_body
 from open_webui.utils.plugin import load_function_module_by_id
-from open_webui.utils.response import normalize_usage
+from open_webui.utils.response import merge_usage, normalize_usage
 from open_webui.utils.sanitize import sanitize_code
 from open_webui.utils.task import (
     get_task_model_id,
@@ -4142,8 +4142,8 @@ async def flush_pending_delta_data(threshold: int = 0):
 
                                         # Normalize and capture usage for DB persistence
                                         if response_metadata.get('usage'):
-                                            response_metadata['usage'] = normalize_usage(response_metadata['usage'])
-                                            usage = response_metadata['usage']
+                                            usage = merge_usage(usage, response_metadata['usage'])
+                                            response_metadata['usage'] = usage
 
                                         processed_data.update(response_metadata)
                                         processed_data.pop('done', None)
@@ -4162,7 +4162,7 @@ async def flush_pending_delta_data(threshold: int = 0):
                                     raw_usage = data.get('usage', {}) or {}
                                     raw_usage.update(data.get('timings', {}))  # llama.cpp
                                     if raw_usage:
-                                        usage = normalize_usage(raw_usage)
+                                        usage = merge_usage(usage, raw_usage)
                                         await event_emitter(
                                             {
                                                 'type': 'chat:completion',
@@ -4260,9 +4260,9 @@ async def flush_pending_delta_data(threshold: int = 0):
                                                         current_response_tool_call['function']['name'] = delta_name
 
                                                     if delta_arguments:
-                                                        current_response_tool_call['function']['arguments'] += (
-                                                            delta_arguments
-                                                        )
+                                                        current_response_tool_call['function'][
+                                                            'arguments'
+                                                        ] += delta_arguments
 
                                         # Emit pending tool calls in real-time
                                         if response_tool_calls:

backend/open_webui/utils/response.py

@@ -1,4 +1,5 @@
 import json
+from numbers import Number
 from uuid import uuid4
 
 from open_webui.utils.misc import (
@@ -50,6 +51,91 @@ def normalize_usage(usage: dict) -> dict:
     return result
 
 
+USAGE_TOKEN_KEYS = {
+    'input_tokens',
+    'output_tokens',
+    'total_tokens',
+    'prompt_tokens',
+    'completion_tokens',
+}
+
+USAGE_COST_KEYS = {
+    'cost',
+    'total_cost',
+    'input_cost',
+    'output_cost',
+    'prompt_cost',
+    'completion_cost',
+}
+
+USAGE_DETAIL_KEYS = {
+    'prompt_tokens_details',
+    'completion_tokens_details',
+    'input_tokens_details',
+    'output_tokens_details',
+}
+
+
+def _is_numeric_usage_value(value) -> bool:
+    return isinstance(value, Number) and not isinstance(value, bool)
+
+
+def _merge_numeric_usage_map(current: dict | None, incoming: dict | None) -> dict:
+    current = current or {}
+    incoming = incoming or {}
+    result = {**current, **incoming}
+
+    for key in set(current) | set(incoming):
+        current_value = current.get(key, 0)
+        incoming_value = incoming.get(key, 0)
+        if isinstance(current_value, dict) or isinstance(incoming_value, dict):
+            result[key] = _merge_numeric_usage_map(
+                current_value if isinstance(current_value, dict) else {},
+                incoming_value if isinstance(incoming_value, dict) else {},
+            )
+        elif _is_numeric_usage_value(current_value) or _is_numeric_usage_value(incoming_value):
+            result[key] = (current_value if _is_numeric_usage_value(current_value) else 0) + (
+                incoming_value if _is_numeric_usage_value(incoming_value) else 0
+            )
+
+    return result
+
+
+def merge_usage(current: dict | None, incoming: dict | None) -> dict:
+    """
+    Merge usage payloads from multiple model calls into one cumulative usage dict.
+
+    Token fields are additive; non-numeric metadata keeps the latest provider value.
+    """
+    current_usage = normalize_usage(current or {}) if current else {}
+    incoming_usage = normalize_usage(incoming or {}) if incoming else {}
+
+    if not incoming_usage:
+        return current_usage
+    if not current_usage:
+        return incoming_usage
+
+    result = {**current_usage, **incoming_usage}
+
+    for key in USAGE_TOKEN_KEYS | USAGE_COST_KEYS:
+        if key in current_usage or key in incoming_usage:
+            current_value = current_usage.get(key, 0)
+            incoming_value = incoming_usage.get(key, 0)
+            if _is_numeric_usage_value(current_value) or _is_numeric_usage_value(incoming_value):
+                result[key] = (current_value if _is_numeric_usage_value(current_value) else 0) + (
+                    incoming_value if _is_numeric_usage_value(incoming_value) else 0
+                )
+
+    for key in USAGE_DETAIL_KEYS:
+        if isinstance(current_usage.get(key), dict) or isinstance(incoming_usage.get(key), dict):
+            result[key] = _merge_numeric_usage_map(
+                current_usage.get(key) if isinstance(current_usage.get(key), dict) else {},
+                incoming_usage.get(key) if isinstance(incoming_usage.get(key), dict) else {},
+            )
+
+    return result
+
+
 def convert_ollama_tool_call_to_openai(tool_calls: list) -> list:
     openai_tool_calls = []
     for tool_call in tool_calls: