refac

backend/open_webui/routers/chats.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from typing import Optional
-
+import asyncio
 
 from open_webui.utils.misc import get_message_list
 from open_webui.socket.main import get_event_emitter
@@ -209,43 +209,12 @@ class ChatStatsExportList(BaseModel):
     page: int
 
 
-@router.get("/stats/export", response_model=ChatStatsExportList)
-def export_chat_stats(
-    request: Request,
-    chat_id: Optional[str] = None,
-    start_time: Optional[int] = None,
-    end_time: Optional[int] = None,
-    page: Optional[int] = 1,
-    user=Depends(get_verified_user),
-):
-    # Check if the user has permission to share/export chats
-    if (user.role != "admin") and (
-        not request.app.state.config.ENABLE_COMMUNITY_SHARING
-    ):
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
-        )
-
-    try:
-        limit = 10  # Fixed limit for export
-        skip = (page - 1) * limit
-
-        # Fetch chats with date filtering
-        filter = {"order_by": "created_at", "direction": "asc"}
-
-        if chat_id:
-            chat = Chats.get_chat_by_id(chat_id)
-            if chat:
-                filter["start_time"] = chat.created_at
-
-        if start_time:
-            filter["start_time"] = start_time
-        if end_time:
-            filter["end_time"] = end_time
+def calculate_chat_stats(user_id, skip=0, limit=10, filter=None):
+    if filter is None:
+        filter = {}
 
     result = Chats.get_chats_by_user_id(
-            user.id,
+        user_id,
         skip=skip,
         limit=limit,
         filter=filter,
@@ -270,7 +239,9 @@ def export_chat_stats(
                     if isinstance(content, str):
                         content_length = len(content)
                     else:
-                            content_length = 0  # Handle cases where content might be None or not string
+                        content_length = (
+                            0  # Handle cases where content might be None or not string
+                        )
 
                     # Extract rating safely
                     rating = message.get("annotation", {}).get("rating")
@@ -370,9 +341,7 @@ def export_chat_stats(
 
             # Construct Chat Body
             chat_body = ChatBody(
-                    history=ChatHistoryStats(
-                        messages=export_messages, currentId=message_id
-                    )
+                history=ChatHistoryStats(messages=export_messages, currentId=message_id)
             )
 
             chat_stat = ChatStatsExport(
@@ -390,10 +359,50 @@ def export_chat_stats(
             log.debug(f"Error exporting stats for chat {chat.id}: {e}")
             continue
 
-        return ChatStatsExportList(
-            items=chat_stats_export_list, total=result.total, page=page
+    return chat_stats_export_list, result.total
+
+
+@router.get("/stats/export", response_model=ChatStatsExportList)
+async def export_chat_stats(
+    request: Request,
+    chat_id: Optional[str] = None,
+    start_time: Optional[int] = None,
+    end_time: Optional[int] = None,
+    page: Optional[int] = 1,
+    user=Depends(get_verified_user),
+):
+    # Check if the user has permission to share/export chats
+    if (user.role != "admin") and (
+        not request.app.state.config.ENABLE_COMMUNITY_SHARING
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
         )
 
+    try:
+        limit = 10  # Fixed limit for export
+        skip = (page - 1) * limit
+
+        # Fetch chats with date filtering
+        filter = {"order_by": "created_at", "direction": "asc"}
+
+        if chat_id:
+            chat = Chats.get_chat_by_id(chat_id)
+            if chat:
+                filter["start_time"] = chat.created_at
+
+        if start_time:
+            filter["start_time"] = start_time
+        if end_time:
+            filter["end_time"] = end_time
+
+        chat_stats_export_list, total = await asyncio.to_thread(
+            calculate_chat_stats, user.id, skip, limit, filter
+        )
+
+        return ChatStatsExportList(items=chat_stats_export_list, total=total, page=page)
+
     except Exception as e:
         log.debug(f"Error exporting chat stats: {e}")
         raise HTTPException(