bhetherman/open-webui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-05-05 04:21:23 +09:00
parent 989d5fd4e2
commit a32d26e61d
1 changed files with 21 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/models/chat_messages.py
+21 -28
View File
@@ -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,