refac
Some checks failed
Create and publish Docker images with specific build args / build-main-image (linux/amd64, ubuntu-latest) (push) Has been cancelled
Create and publish Docker images with specific build args / build-main-image (linux/arm64, ubuntu-24.04-arm) (push) Has been cancelled
Create and publish Docker images with specific build args / build-cuda-image (linux/amd64, ubuntu-latest) (push) Has been cancelled
Create and publish Docker images with specific build args / build-cuda-image (linux/arm64, ubuntu-24.04-arm) (push) Has been cancelled
Create and publish Docker images with specific build args / build-cuda126-image (linux/amd64, ubuntu-latest) (push) Has been cancelled
Create and publish Docker images with specific build args / build-cuda126-image (linux/arm64, ubuntu-24.04-arm) (push) Has been cancelled
Create and publish Docker images with specific build args / build-ollama-image (linux/amd64, ubuntu-latest) (push) Has been cancelled
Create and publish Docker images with specific build args / build-ollama-image (linux/arm64, ubuntu-24.04-arm) (push) Has been cancelled
Create and publish Docker images with specific build args / build-slim-image (linux/amd64, ubuntu-latest) (push) Has been cancelled
Create and publish Docker images with specific build args / build-slim-image (linux/arm64, ubuntu-24.04-arm) (push) Has been cancelled
Python CI / Format Backend (push) Has been cancelled
Frontend Build / Format & Build Frontend (push) Has been cancelled
Frontend Build / Frontend Unit Tests (push) Has been cancelled
Create and publish Docker images with specific build args / merge-main-images (push) Has been cancelled
Create and publish Docker images with specific build args / merge-cuda-images (push) Has been cancelled
Create and publish Docker images with specific build args / merge-cuda126-images (push) Has been cancelled
Create and publish Docker images with specific build args / merge-ollama-images (push) Has been cancelled
Create and publish Docker images with specific build args / merge-slim-images (push) Has been cancelled
Create and publish Docker images with specific build args / copy-to-dockerhub (, main) (push) Has been cancelled
Create and publish Docker images with specific build args / copy-to-dockerhub (-cuda, cuda) (push) Has been cancelled
Create and publish Docker images with specific build args / copy-to-dockerhub (-cuda126, cuda126) (push) Has been cancelled
Create and publish Docker images with specific build args / copy-to-dockerhub (-ollama, ollama) (push) Has been cancelled
Create and publish Docker images with specific build args / copy-to-dockerhub (-slim, slim) (push) Has been cancelled

This commit is contained in:
Timothy Jaeryang Baek 2026-05-05 04:21:23 +09:00
parent 989d5fd4e2
commit a32d26e61d

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,