From a32d26e61d24d9f63650faed5cb8909ed90af661 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 5 May 2026 04:21:23 +0900 Subject: [PATCH] refac --- backend/open_webui/models/chat_messages.py | 49 ++++++++++------------ 1 file changed, 21 insertions(+), 28 deletions(-) diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py index 0758e0354d..b1768d0ff2 100644 --- a/backend/open_webui/models/chat_messages.py +++ b/backend/open_webui/models/chat_messages.py @@ -48,6 +48,25 @@ def get_usage(data: dict) -> Optional[dict]: return normalize_usage(usage) if usage else None +def _token_columns(dialect: str): + """Return (input_tokens, output_tokens) SQL column expressions. + + Falls back to OpenAI-style keys (prompt_tokens / completion_tokens) + when the normalized keys are absent. + """ + if dialect == 'sqlite': + extract = lambda key: cast(func.json_extract(ChatMessage.usage, f'$.{key}'), Integer) + elif dialect == 'postgresql': + extract = lambda key: cast(func.json_extract_path_text(ChatMessage.usage, key), Integer) + else: + raise NotImplementedError(f'Unsupported dialect: {dialect}') + + return ( + func.coalesce(extract('input_tokens'), extract('prompt_tokens')), + func.coalesce(extract('output_tokens'), extract('completion_tokens')), + ) + + #################### # ChatMessage DB Schema #################### @@ -343,20 +362,7 @@ class ChatMessageTable: bind = await db.connection() dialect = bind.dialect.name - if dialect == 'sqlite': - input_tokens = cast(func.json_extract(ChatMessage.usage, '$.input_tokens'), Integer) - output_tokens = cast(func.json_extract(ChatMessage.usage, '$.output_tokens'), Integer) - elif dialect == 'postgresql': - input_tokens = cast( - func.json_extract_path_text(ChatMessage.usage, 'input_tokens'), - Integer, - ) - output_tokens = cast( - func.json_extract_path_text(ChatMessage.usage, 'output_tokens'), - Integer, - ) - else: - raise NotImplementedError(f'Unsupported dialect: {dialect}') + input_tokens, output_tokens = _token_columns(dialect) stmt = select( ChatMessage.model_id, @@ -404,20 +410,7 @@ class ChatMessageTable: bind = await db.connection() dialect = bind.dialect.name - if dialect == 'sqlite': - input_tokens = cast(func.json_extract(ChatMessage.usage, '$.input_tokens'), Integer) - output_tokens = cast(func.json_extract(ChatMessage.usage, '$.output_tokens'), Integer) - elif dialect == 'postgresql': - input_tokens = cast( - func.json_extract_path_text(ChatMessage.usage, 'input_tokens'), - Integer, - ) - output_tokens = cast( - func.json_extract_path_text(ChatMessage.usage, 'output_tokens'), - Integer, - ) - else: - raise NotImplementedError(f'Unsupported dialect: {dialect}') + input_tokens, output_tokens = _token_columns(dialect) stmt = select( ChatMessage.user_id,