refac

backend/open_webui/models/chat_messages.py

@@ -1,8 +1,10 @@
 from contextlib import nullcontext
+import os
 import time
 import uuid
 from collections import Counter
 from datetime import datetime, timedelta
+from pathlib import Path
 from typing import Any, Optional
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
@@ -17,10 +19,12 @@ from sqlalchemy import (
     Boolean,
     Column,
     ForeignKey,
+    Float,
     Index,
     Integer,
     Text,
     cast,
+    case,
     delete,
     func,
     literal,
@@ -89,6 +93,36 @@ def _token_columns(dialect: str):
     )
 
 
+def _daily_usage_bucket(dialect: str, bucket_offset: int):
+    """Match Python's server-local day and legacy timestamp normalization in SQL."""
+    if dialect not in {'sqlite', 'postgresql'}:
+        return None
+    if dialect == 'postgresql':
+        # The database server/session can use a different timezone from the application.
+        zone = os.environ.get('TZ')
+        if zone is None:
+            localtime = str(Path('/etc/localtime').resolve())
+            if '/zoneinfo/' in localtime:
+                zone = localtime.split('/zoneinfo/', 1)[1]
+            elif time.timezone == 0 and not time.daylight:
+                zone = 'UTC'
+        try:
+            ZoneInfo(zone or '')
+        except (ValueError, ZoneInfoNotFoundError):
+            # Retain Python bucketing when the OS timezone has no portable IANA name.
+            return None
+
+    now = time.time()
+    seconds = case(
+        (ChatMessage.created_at > 10_000_000_000, cast(ChatMessage.created_at, Float) / 1000),
+        else_=ChatMessage.created_at,
+    )
+    timestamp = case((seconds.between(1577836800, now + 86400), seconds), else_=now) + bucket_offset
+    if dialect == 'sqlite':
+        return func.strftime('%Y-%m-%d', timestamp, 'unixepoch', 'localtime')
+    return func.to_char(func.timezone(zone, func.to_timestamp(timestamp)), 'YYYY-MM-DD')
+
+
 def _extract_tool_names(value: Any) -> list[str]:
     names: list[str] = []
 
@@ -520,6 +554,37 @@ class ChatMessageTable:
             return True
 
     # Analytics methods
+    async def get_analytics_summary(
+        self,
+        start_date: Optional[int] = None,
+        end_date: Optional[int] = None,
+        group_id: Optional[str] = None,
+        db: Optional[AsyncSession] = None,
+    ) -> dict[str, int]:
+        """Return dashboard totals without materializing every model, user, and chat group."""
+        async with get_async_db_context(db) as db:
+            from open_webui.models.groups import group_user_memberships
+
+            stmt = select(
+                func.count(ChatMessage.model_id).label('total_messages'),
+                func.count(distinct(ChatMessage.chat_id)).label('total_chats'),
+                func.count(distinct(ChatMessage.model_id)).label('total_models'),
+                # GROUP BY user_id included the anonymous/null group in the old summary.
+                (
+                    func.count(distinct(ChatMessage.user_id))
+                    + case((func.count() > func.count(ChatMessage.user_id), 1), else_=0)
+                ).label('total_users'),
+            ).where(ChatMessage.role == 'assistant')
+            if start_date:
+                stmt = stmt.where(ChatMessage.created_at >= start_date)
+            if end_date:
+                stmt = stmt.where(ChatMessage.created_at <= end_date)
+            if group_id:
+                group_users = select(group_user_memberships([group_id], True).c.user_id).scalar_subquery()
+                stmt = stmt.where(ChatMessage.user_id.in_(group_users))
+
+            return dict((await db.execute(stmt)).mappings().one())
+
     async def get_message_count_by_model(
         self,
         start_date: Optional[int] = None,
@@ -547,19 +612,20 @@ class ChatMessageTable:
             result = await db.execute(stmt)
             return {row.model_id: row.count for row in result.all()}
 
-    async def get_unique_counts_by_model(
+    async def get_model_analytics(
         self,
         start_date: Optional[int] = None,
         end_date: Optional[int] = None,
         group_id: Optional[str] = None,
         db: Optional[AsyncSession] = None,
     ) -> dict[str, dict]:
-        """Count distinct users and chats per model."""
+        """Count messages, distinct users, and chats per model in one pass."""
         async with get_async_db_context(db) as db:
             from open_webui.models.groups import group_user_memberships
 
             stmt = select(
                 ChatMessage.model_id,
+                func.count(ChatMessage.id).label('count'),
                 func.count(distinct(ChatMessage.user_id)).label('unique_users'),
                 func.count(distinct(ChatMessage.chat_id)).label('unique_chats'),
             ).filter(
@@ -579,6 +645,7 @@ class ChatMessageTable:
             result = await db.execute(stmt)
             return {
                 row.model_id: {
+                    'count': row.count,
                     'unique_users': row.unique_users,
                     'unique_chats': row.unique_chats,
                 }
@@ -635,14 +702,14 @@ class ChatMessageTable:
                 for row in result.all()
             }
 
-    async def get_token_usage_by_user(
+    async def get_user_analytics(
         self,
         start_date: Optional[int] = None,
         end_date: Optional[int] = None,
         group_id: Optional[str] = None,
         db: Optional[AsyncSession] = None,
     ) -> dict[str, dict]:
-        """Aggregate token usage by user using database-level aggregation."""
+        """Aggregate all message counts and token usage per user in one pass."""
         async with get_async_db_context(db) as db:
             from open_webui.models.groups import group_user_memberships
 
@@ -651,6 +718,10 @@ class ChatMessageTable:
 
             input_tokens, output_tokens = _token_columns(dialect)
 
+            # The previous token query excluded anonymous users; keep those totals at zero.
+            input_tokens = case((ChatMessage.user_id.isnot(None), input_tokens), else_=0)
+            output_tokens = case((ChatMessage.user_id.isnot(None), output_tokens), else_=0)
+
             stmt = select(
                 ChatMessage.user_id,
                 func.coalesce(func.sum(input_tokens), 0).label('input_tokens'),
@@ -658,8 +729,6 @@ class ChatMessageTable:
                 func.count(ChatMessage.id).label('message_count'),
             ).filter(
                 ChatMessage.role == 'assistant',
-                ChatMessage.user_id.isnot(None),
-                ChatMessage.usage.isnot(None),
             )
 
             if start_date:
@@ -987,13 +1056,18 @@ class ChatMessageTable:
         async with get_async_db_context(db) as db:
             from open_webui.models.groups import group_user_memberships
 
+            bind = await db.connection()
+            day = _daily_usage_bucket(bind.dialect.name, bucket_offset) if granularity == 'daily' else None
             value = literal(1)
             if metric == 'tokens':
-                bind = await db.connection()
                 input_tokens, output_tokens = _token_columns(bind.dialect.name)
                 value = func.coalesce(input_tokens, 0) + func.coalesce(output_tokens, 0)
 
-            stmt = select(ChatMessage.created_at, ChatMessage.model_id, value).filter(
+            stmt = select(
+                day if day is not None else ChatMessage.created_at,
+                ChatMessage.model_id,
+                func.sum(value) if day is not None else value,
+            ).filter(
                 ChatMessage.role == 'assistant',
                 ChatMessage.model_id.isnot(None),
             )
@@ -1006,10 +1080,12 @@ class ChatMessageTable:
                 group_users = select(group_user_memberships([group_id], True).c.user_id).scalar_subquery()
                 stmt = stmt.filter(ChatMessage.user_id.in_(group_users))
 
+            if day is not None:
+                stmt = stmt.group_by(day, ChatMessage.model_id)
             result = await db.execute(stmt)
             counts: dict[str, dict[str, int]] = {}
             for timestamp, model_id, value in result.all():
-                date_str = bucket_start(timestamp).strftime(date_format)
+                date_str = timestamp if day is not None else bucket_start(timestamp).strftime(date_format)
                 bucket = counts.setdefault(date_str, {})
                 bucket[model_id] = bucket.get(model_id, 0) + int(value)
 

backend/open_webui/routers/analytics.py

@@ -64,20 +64,12 @@ async def get_model_analytics(
     db: AsyncSession = Depends(get_async_session),
 ):
     """Get message counts per model."""
-    counts = await ChatMessages.get_message_count_by_model(
-        start_date=start_date, end_date=end_date, group_id=group_id, db=db
-    )
-    unique_counts = await ChatMessages.get_unique_counts_by_model(
+    counts = await ChatMessages.get_model_analytics(
         start_date=start_date, end_date=end_date, group_id=group_id, db=db
     )
     models = [
-        ModelAnalyticsEntry(
-            model_id=model_id,
-            count=count,
-            unique_users=unique_counts.get(model_id, {}).get('unique_users', 0),
-            unique_chats=unique_counts.get(model_id, {}).get('unique_chats', 0),
-        )
-        for model_id, count in sorted(counts.items(), key=lambda x: -x[1])
+        ModelAnalyticsEntry(model_id=model_id, **data)
+        for model_id, data in sorted(counts.items(), key=lambda x: -x[1]['count'])
     ]
     return ModelAnalyticsResponse(models=models)
 
@@ -93,19 +85,16 @@ async def get_user_analytics(
     db: AsyncSession = Depends(get_async_session),
 ):
     """Get message counts and token usage per user with user info."""
-    counts = await ChatMessages.get_message_count_by_user(
-        start_date=start_date, end_date=end_date, group_id=group_id, db=db
-    )
-    token_usage = await ChatMessages.get_token_usage_by_user(
+    usage = await ChatMessages.get_user_analytics(
         start_date=start_date, end_date=end_date, group_id=group_id, db=db
     )
 
     # Rank by the selected metric before limiting the leaderboard.
     top_user_ids = sorted(
-        counts,
+        usage,
         key=lambda uid: (
-            -token_usage.get(uid, {}).get('total_tokens', 0) if metric == 'tokens' else -counts[uid],
-            -counts[uid],
+            -usage[uid]['total_tokens'] if metric == 'tokens' else -usage[uid]['message_count'],
+            -usage[uid]['message_count'],
         ),
     )[:limit]
     user_info = {u.id: u for u in await Users.get_users_by_user_ids(top_user_ids, db=db)}
@@ -113,13 +102,13 @@ async def get_user_analytics(
     users = []
     for user_id in top_user_ids:
         u = user_info.get(user_id)
-        tokens = token_usage.get(user_id, {})
+        tokens = usage[user_id]
         users.append(
             UserAnalyticsEntry(
                 user_id=user_id,
                 name=u.name if u else None,
                 email=u.email if u else None,
-                count=counts[user_id],
+                count=tokens['message_count'],
                 input_tokens=tokens.get('input_tokens', 0),
                 output_tokens=tokens.get('output_tokens', 0),
                 total_tokens=tokens.get('total_tokens', 0),
@@ -176,22 +165,11 @@ async def get_summary(
     db: AsyncSession = Depends(get_async_session),
 ):
     """Get summary statistics for the dashboard."""
-    model_counts = await ChatMessages.get_message_count_by_model(
-        start_date=start_date, end_date=end_date, group_id=group_id, db=db
-    )
-    user_counts = await ChatMessages.get_message_count_by_user(
-        start_date=start_date, end_date=end_date, group_id=group_id, db=db
-    )
-    chat_counts = await ChatMessages.get_message_count_by_chat(
+    summary = await ChatMessages.get_analytics_summary(
         start_date=start_date, end_date=end_date, group_id=group_id, db=db
     )
 
-    return SummaryResponse(
-        total_messages=sum(model_counts.values()),
-        total_chats=len(chat_counts),
-        total_models=len(model_counts),
-        total_users=len(user_counts),
-    )
+    return SummaryResponse(**summary)
 
 
 class ComparisonStatsEntry(BaseModel):