This commit is contained in:
Timothy Jaeryang Baek
2026-04-02 08:34:49 -05:00
parent 2863e9f8c4
commit 4dea4fdf54

View File

@@ -5,6 +5,7 @@ from typing import Any, Optional
from sqlalchemy.orm import Session
from open_webui.internal.db import Base, get_db_context
from open_webui.utils.response import normalize_usage
from pydantic import BaseModel, ConfigDict
from sqlalchemy import (
@@ -41,6 +42,12 @@ def _normalize_timestamp(timestamp: int) -> float:
return timestamp
def get_usage(data: dict) -> Optional[dict]:
"""Extract and normalize usage from message data."""
usage = data.get('usage') or (data.get('info') or {}).get('usage')
return normalize_usage(usage) if usage else None
####################
# ChatMessage DB Schema
####################
@@ -163,11 +170,8 @@ class ChatMessageTable:
existing.status_history = data.get('status_history') or data.get('statusHistory')
if 'error' in data:
existing.error = data.get('error')
# Extract usage - check direct field first, then info.usage
usage = data.get('usage')
if not usage:
info = data.get('info', {})
usage = info.get('usage') if info else None
# 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)
@@ -179,11 +183,8 @@ class ChatMessageTable:
return ChatMessageModel.model_validate(existing)
else:
# Insert new
# Extract usage - check direct field first, then info.usage
usage = data.get('usage')
if not usage:
info = data.get('info', {})
usage = info.get('usage') if info else None
# Extract and normalize usage
usage = get_usage(data)
message = ChatMessage(
id=composite_id,
chat_id=chat_id,