diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py index e0f00da56..f3058f4ff 100644 --- a/backend/open_webui/models/chat_messages.py +++ b/backend/open_webui/models/chat_messages.py @@ -1,3 +1,4 @@ +import json import time import uuid from typing import Any, Optional @@ -279,11 +280,34 @@ class ChatMessageTable: end_date: Optional[int] = None, db: Optional[Session] = None, ) -> dict[str, dict]: - """Aggregate token usage by model. Works with SQLite and PostgreSQL.""" + """Aggregate token usage by model using database-level aggregation.""" with get_db_context(db) as db: + from sqlalchemy import func, cast, Integer + + dialect = db.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( + ChatMessage.usage["input_tokens"].astext, Integer + ) + output_tokens = cast( + ChatMessage.usage["output_tokens"].astext, Integer + ) + else: + raise NotImplementedError(f"Unsupported dialect: {dialect}") + query = db.query( ChatMessage.model_id, - ChatMessage.usage, + func.coalesce(func.sum(input_tokens), 0).label("input_tokens"), + func.coalesce(func.sum(output_tokens), 0).label("output_tokens"), + func.count(ChatMessage.id).label("message_count"), ).filter( ChatMessage.role == "assistant", ChatMessage.model_id.isnot(None), @@ -295,27 +319,17 @@ class ChatMessageTable: if end_date: query = query.filter(ChatMessage.created_at <= end_date) - results = query.all() + results = query.group_by(ChatMessage.model_id).all() - # Aggregate in Python for cross-database compatibility - usage_by_model: dict[str, dict] = {} - for model_id, usage in results: - if model_id not in usage_by_model: - usage_by_model[model_id] = { - "input_tokens": 0, - "output_tokens": 0, - "message_count": 0, - } - - usage_by_model[model_id]["input_tokens"] += usage.get("input_tokens") or 0 - usage_by_model[model_id]["output_tokens"] += usage.get("output_tokens") or 0 - usage_by_model[model_id]["message_count"] += 1 - - # Add total_tokens - for data in usage_by_model.values(): - data["total_tokens"] = data["input_tokens"] + data["output_tokens"] - - return usage_by_model + return { + row.model_id: { + "input_tokens": row.input_tokens, + "output_tokens": row.output_tokens, + "total_tokens": row.input_tokens + row.output_tokens, + "message_count": row.message_count, + } + for row in results + } def get_message_count_by_user( self, diff --git a/backend/open_webui/routers/analytics.py b/backend/open_webui/routers/analytics.py index d71da0f4f..db66e9943 100644 --- a/backend/open_webui/routers/analytics.py +++ b/backend/open_webui/routers/analytics.py @@ -192,3 +192,46 @@ async def get_daily_stats( for date, models in sorted(counts.items()) ] ) + + +class TokenUsageEntry(BaseModel): + model_id: str + input_tokens: int + output_tokens: int + total_tokens: int + message_count: int + + +class TokenUsageResponse(BaseModel): + models: list[TokenUsageEntry] + total_input_tokens: int + total_output_tokens: int + total_tokens: int + + +@router.get("/tokens", response_model=TokenUsageResponse) +async def get_token_usage( + start_date: Optional[int] = Query(None), + end_date: Optional[int] = Query(None), + user=Depends(get_admin_user), + db: Session = Depends(get_session), +): + """Get token usage aggregated by model.""" + usage = ChatMessages.get_token_usage_by_model( + start_date=start_date, end_date=end_date, db=db + ) + + models = [ + TokenUsageEntry(model_id=model_id, **data) + for model_id, data in sorted(usage.items(), key=lambda x: -x[1]["total_tokens"]) + ] + + total_input = sum(m.input_tokens for m in models) + total_output = sum(m.output_tokens for m in models) + + return TokenUsageResponse( + models=models, + total_input_tokens=total_input, + total_output_tokens=total_output, + total_tokens=total_input + total_output, + ) diff --git a/src/lib/apis/analytics/index.ts b/src/lib/apis/analytics/index.ts index 3aeb71139..19ef2a388 100644 --- a/src/lib/apis/analytics/index.ts +++ b/src/lib/apis/analytics/index.ts @@ -193,3 +193,39 @@ export const getDailyStats = async ( return res; }; + +export const getTokenUsage = async ( + token: string = '', + startDate: number | null = null, + endDate: number | null = null +) => { + let error = null; + + const searchParams = new URLSearchParams(); + if (startDate) searchParams.append('start_date', startDate.toString()); + if (endDate) searchParams.append('end_date', endDate.toString()); + + const res = await fetch(`${WEBUI_API_BASE_URL}/analytics/tokens?${searchParams.toString()}`, { + method: 'GET', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + authorization: `Bearer ${token}` + } + }) + .then(async (res) => { + if (!res.ok) throw await res.json(); + return res.json(); + }) + .catch((err) => { + error = err.detail; + console.error(err); + return null; + }); + + if (error) { + throw error; + } + + return res; +}; diff --git a/src/lib/components/admin/Analytics/Dashboard.svelte b/src/lib/components/admin/Analytics/Dashboard.svelte index 5efc41b55..e49275690 100644 --- a/src/lib/components/admin/Analytics/Dashboard.svelte +++ b/src/lib/components/admin/Analytics/Dashboard.svelte @@ -1,12 +1,13 @@