From 95391221dfabbfcd9090ab472b1c02a4c75c0387 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 23 Jun 2026 23:35:44 +0200 Subject: [PATCH] refac --- backend/open_webui/models/chat_messages.py | 8 +- backend/open_webui/utils/middleware.py | 14 ++-- backend/open_webui/utils/response.py | 86 ++++++++++++++++++++++ 3 files changed, 96 insertions(+), 12 deletions(-) diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py index 8bb170578a..5579bdbe95 100644 --- a/backend/open_webui/models/chat_messages.py +++ b/backend/open_webui/models/chat_messages.py @@ -6,7 +6,7 @@ from typing import Any, Optional 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 @@ class ChatMessageTable: # 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) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 82ca932646..4c17d720a5 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -112,7 +112,7 @@ from open_webui.utils.misc import ( ) 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 streaming_chat_response_handler(response, ctx): # 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 streaming_chat_response_handler(response, ctx): 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 streaming_chat_response_handler(response, ctx): 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: diff --git a/backend/open_webui/utils/response.py b/backend/open_webui/utils/response.py index 7bc5375480..de3ede0a07 100644 --- a/backend/open_webui/utils/response.py +++ b/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: