Update app.py

app.py

@@ -1,67 +1,92 @@
 """
-Chainlit UI for AI Agent with Login/Registration
+Chainlit UI for AI Agent with Login/Registration - FIXED VERSION
 """
 
 import chainlit as cl
-import requests
+import httpx  # Thay thế requests bằng httpx cho async support
 from typing import Optional, Dict
 import json
 import os
+import asyncio
 
 # API Configuration
 API_BASE_URL = os.getenv("API_BASE_URL")
 
+# Timeout settings (giây)
+API_TIMEOUT = 300.0  # 5 phút cho chat API (có thể điều chỉnh)
+AUTH_TIMEOUT = 30.0  # 30 giây cho login/register
+
 # ============================================
-# API HELPER FUNCTIONS
+# API HELPER FUNCTIONS - CONVERTED TO ASYNC
 # ============================================
 
-def check_user_exists(user_id: str) -> bool:
+async def check_user_exists(user_id: str) -> bool:
     """Check if user exists in the system"""
     try:
-        response = requests.get(f"{API_BASE_URL}/user-exists/{user_id}")
-        if response.status_code == 200:
-            return response.json().get("exists", False)
+        async with httpx.AsyncClient(timeout=AUTH_TIMEOUT) as client:
+            response = await client.get(f"{API_BASE_URL}/user-exists/{user_id}")
+            if response.status_code == 200:
+                return response.json().get("exists", False)
+            return False
+    except httpx.TimeoutException:
+        print(f"Timeout checking user: {user_id}")
         return False
     except Exception as e:
         print(f"Error checking user: {e}")
         return False
 
 
-def register_user(user_id: str) -> Dict:
+async def register_user(user_id: str) -> Dict:
     """Register a new user"""
     try:
-        response = requests.post(
-            f"{API_BASE_URL}/register",
-            json={"user_id": user_id}
-        )
-        if response.status_code == 201:
-            return {"success": True, "message": response.json().get("message", "Đăng ký thành công!")}
-        else:
-            error_msg = response.json().get("detail", "Đăng ký thất bại")
-            return {"success": False, "message": error_msg}
+        async with httpx.AsyncClient(timeout=AUTH_TIMEOUT) as client:
+            response = await client.post(
+                f"{API_BASE_URL}/register",
+                json={"user_id": user_id}
+            )
+            if response.status_code == 201:
+                return {"success": True, "message": response.json().get("message", "Đăng ký thành công!")}
+            else:
+                error_msg = response.json().get("detail", "Đăng ký thất bại")
+                return {"success": False, "message": error_msg}
+    except httpx.TimeoutException:
+        return {"success": False, "message": "Timeout: Server không phản hồi"}
+    except httpx.ConnectError:
+        return {"success": False, "message": "Không thể kết nối đến server"}
     except Exception as e:
-        return {"success": False, "message": f"Lỗi kết nối: {str(e)}"}
+        return {"success": False, "message": f"Lỗi: {str(e)}"}
 
 
-def send_chat_message(query: str, user_id: str) -> Dict:
-    """Send chat message to API"""
+async def send_chat_message(query: str, user_id: str) -> Dict:
+    """Send chat message to API with streaming support"""
     try:
-        response = requests.post(
-            f"{API_BASE_URL}/chat",
-            json={"query": query, "user_id": user_id}
-        )
-        if response.status_code == 200:
-            data = response.json()
-            return {
-                "success": True,
-                "answer": data.get("answer", ""),
-                "references": data.get("references", "Không có tài liệu tham khảo")
-            }
-        else:
-            error_msg = response.json().get("detail", "Lỗi xử lý câu hỏi")
-            return {"success": False, "message": error_msg}
+        # Sử dụng timeout dài hơn cho chat vì có thể xử lý lâu
+        async with httpx.AsyncClient(timeout=API_TIMEOUT) as client:
+            response = await client.post(
+                f"{API_BASE_URL}/chat",
+                json={"query": query, "user_id": user_id}
+            )
+            
+            if response.status_code == 200:
+                data = response.json()
+                return {
+                    "success": True,
+                    "answer": data.get("answer", ""),
+                    "references": data.get("references", "Không có tài liệu tham khảo")
+                }
+            else:
+                error_msg = response.json().get("detail", "Lỗi xử lý câu hỏi")
+                return {"success": False, "message": error_msg}
+                
+    except httpx.TimeoutException:
+        return {
+            "success": False, 
+            "message": f"⏱️ Server đang xử lý quá lâu (>{API_TIMEOUT}s). Vui lòng thử lại hoặc đặt câu hỏi đơn giản hơn."
+        }
+    except httpx.ConnectError:
+        return {"success": False, "message": "❌ Không thể kết nối đến server. Vui lòng kiểm tra API_BASE_URL."}
     except Exception as e:
-        return {"success": False, "message": f"Lỗi kết nối API: {str(e)}"}
+        return {"success": False, "message": f"❌ Lỗi: {str(e)}"}
 
 
 # ============================================
@@ -171,12 +196,14 @@ async def handle_command(command: str):
         loading_msg = cl.Message(content="🔍 Đang kiểm tra tài khoản...", author="Kumiko")
         await loading_msg.send()
         
-        # Check if user exists
-        if check_user_exists(user_id):
-            await loading_msg.remove()
+        # Check if user exists (now async)
+        exists = await check_user_exists(user_id)
+        
+        await loading_msg.remove()
+        
+        if exists:
             await show_chat_interface(user_id)
         else:
-            await loading_msg.remove()
             await cl.Message(
                 content=f"❌ User ID `{user_id}` không tồn tại!\n\nVui lòng đăng ký bằng: `/register {user_id}`",
                 author="Kumiko"
@@ -198,8 +225,9 @@ async def handle_command(command: str):
         loading_msg = cl.Message(content="📝 Đang đăng ký tài khoản...", author="Kumiko")
         await loading_msg.send()
         
-        # Register user
-        result = register_user(user_id)
+        # Register user (now async)
+        result = await register_user(user_id)
+        
         await loading_msg.remove()
         
         if result["success"]:
@@ -238,44 +266,55 @@ Sau khi đăng nhập, bạn có thể chat bình thường. Hệ thống sẽ t
 
 
 async def process_chat_message(query: str, user_id: str):
-    """Process chat message and display response"""
+    """Process chat message and display response with progress updates"""
     # Show thinking message
     thinking_msg = cl.Message(content="🤔 Đang suy nghĩ...", author="Kumiko")
     await thinking_msg.send()
     
-    # Send request to API
-    result = send_chat_message(query, user_id)
-    
-    # Remove thinking message
-    await thinking_msg.remove()
-    
-    if not result["success"]:
-        await cl.Message(content=f"❌ Lỗi: {result['message']}", author="Kumiko").send()
-        return
-    
-    answer = result["answer"]
-    references = result["references"]
-    
-    # Store references in session for the button
-    cl.user_session.set("last_references", references)
-    
-    # Create message with action button
-    actions = [
-        cl.Action(
-            name="show_references",
-            value="show",
-            label="𝄢 Tài liệu tham khảo",
-            description="Xem các nguồn tham khảo",
-            payload={"references": references}
-        )
-    ]
-    
-    # Send answer with reference button
-    await cl.Message(
-        content=answer,
-        actions=actions,
-        author="Kumiko"
-    ).send()
+    try:
+        # Tạo task để track progress
+        start_time = asyncio.get_event_loop().time()
+        
+        # Send request to API (now async)
+        result = await send_chat_message(query, user_id)
+        
+        # Remove thinking message
+        await thinking_msg.remove()
+        
+        if not result["success"]:
+            await cl.Message(content=f"{result['message']}", author="Kumiko").send()
+            return
+        
+        answer = result["answer"]
+        references = result["references"]
+        
+        # Store references in session for the button
+        cl.user_session.set("last_references", references)
+        
+        # Create message with action button
+        actions = [
+            cl.Action(
+                name="show_references",
+                value="show",
+                label="𝄢 Tài liệu tham khảo",
+                description="Xem các nguồn tham khảo",
+                payload={"references": references}
+            )
+        ]
+        
+        # Send answer with reference button
+        await cl.Message(
+            content=answer,
+            actions=actions,
+            author="Kumiko"
+        ).send()
+        
+    except Exception as e:
+        await thinking_msg.remove()
+        await cl.Message(
+            content=f"❌ Đã xảy ra lỗi không mong muốn: {str(e)}",
+            author="Kumiko"
+        ).send()
 
 
 @cl.action_callback("show_references")
@@ -299,6 +338,7 @@ async def end():
     """Handle chat end"""
     print("Chat session ended")
 
+
 # ============================================
 # CONFIGURATION
 # ============================================
@@ -321,17 +361,4 @@ async def set_starters():
             label="Trợ giúp",
             message="/help",
         ),
-    ]
-
-
-# Custom CSS (optional - save as .chainlit/config.toml)
-# """
-# [UI]
-# name = "AI Agent Chat"
-# default_collapse_content = true
-# default_expand_messages = false
-# hide_cot = false
-
-# [UI.theme]
-# primary_color = "#2563eb"
-# """
+    ]