Thêm memory theo user id

app.py

@@ -1,199 +1,316 @@
+"""
+Chainlit UI for AI Agent with Login/Registration
+"""
+
 import chainlit as cl
-import aiohttp
+import requests
+from typing import Optional, Dict
+import json
 import asyncio
-from typing import Optional, Dict, Any
 
 # API Configuration
 API_BASE_URL = "https://kumiko-v2.onrender.com"
-API_TIMEOUT = 360
-
-class APIClient:
-    """Simple async API client for Kumiko Agent"""
-    
-    def __init__(self, base_url: str = API_BASE_URL):
-        self.base_url = base_url
-        self.session = None
-    
-    async def _get_session(self):
-        """Get or create aiohttp session"""
-        if self.session is None or self.session.closed:
-            timeout = aiohttp.ClientTimeout(total=API_TIMEOUT)
-            self.session = aiohttp.ClientSession(timeout=timeout)
-        return self.session
-    
-    async def _make_request(self, method: str, endpoint: str, data: Optional[Dict] = None) -> Dict[str, Any]:
-        """Make HTTP request to API"""
-        session = await self._get_session()
-        url = f"{self.base_url}{endpoint}"
-        
-        try:
-            if method.upper() == "GET":
-                async with session.get(url) as response:
-                    return await self._handle_response(response)
-            elif method.upper() == "POST":
-                async with session.post(url, json=data) as response:
-                    return await self._handle_response(response)
-            elif method.upper() == "DELETE":
-                async with session.delete(url) as response:
-                    return await self._handle_response(response)
-        except aiohttp.ClientError as e:
-            raise Exception(f"Không thể kết nối API: {str(e)}")
-    
-    async def _handle_response(self, response: aiohttp.ClientResponse) -> Dict[str, Any]:
-        """Handle API response"""
-        if response.status == 200:
-            return await response.json()
+
+# ============================================
+# API HELPER FUNCTIONS
+# ============================================
+
+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)
+        return False
+    except Exception as e:
+        print(f"Error checking user: {e}")
+        return False
+
+
+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_text = await response.text()
-            raise Exception(f"API error {response.status}: {error_text}")
-    
-    async def health_check(self) -> bool:
-        """Check if API is healthy"""
-        try:
-            result = await self._make_request("GET", "/health")
-            return result.get("status") == "healthy"
-        except:
-            return False
-    
-    async def chat(self, query: str, max_token_limit: int = 2000) -> Dict[str, Any]:
-        """Send chat message to /chat endpoint"""
-        data = {
-            "query": query,
-            "max_token_limit": max_token_limit
-        }
-        return await self._make_request("POST", "/chat", data)
-    
-    async def clear_all(self) -> Dict[str, Any]:
-        """Clear all memory and caches via /clear-all endpoint"""
-        return await self._make_request("DELETE", "/clear-all")
-    
-    async def close(self):
-        """Close the session"""
-        if self.session and not self.session.closed:
-            await self.session.close()
-
-# Global API client
-api_client = APIClient()
+            error_msg = response.json().get("detail", "Đăng ký thất bại")
+            return {"success": False, "message": error_msg}
+    except Exception as e:
+        return {"success": False, "message": f"Lỗi kết nối: {str(e)}"}
+
+
+def send_chat_message(query: str, user_id: str) -> Dict:
+    """Send chat message to API"""
+    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}
+    except Exception as e:
+        return {"success": False, "message": f"Lỗi kết nối API: {str(e)}"}
+
+
+# ============================================
+# CHAINLIT EVENT HANDLERS
+# ============================================
 
 @cl.on_chat_start
 async def start():
     """Initialize chat session"""
-    # Check API health
-    is_healthy = await api_client.health_check()
+    # Check if user is already logged in
+    user_id = cl.user_session.get("user_id")
     
-    if not is_healthy:
-        await cl.Message(
-            content="❌ **Không thể kết nối đến Kumiko API!**\n\n"
-                   "Vui lòng kiểm tra:\n"
-                   "- API server đã được khởi động tại http://127.0.0.1:7861\n"
-                   "- Server hoạt động bình thường",
-            author="Kumiko"
-        ).send()
-        return
+    if not user_id:
+        await show_login_screen()
+    else:
+        await show_chat_interface(user_id)
+
+
+async def show_login_screen():
+    """Display login/registration screen"""
+    welcome_msg = """
+### 🎼 Chào mừng đến với Kumiko v2! 📯
+
+**Đăng nhập** hoặc **đăng ký** để bắt đầu trò chuyện với mình nhé!
+
+---
+
+#### 📝 Hướng dẫn:
+- **Đăng nhập**: Gõ `/login <user_id>` (ví dụ: `/login john`)
+- **Đăng ký**: Gõ `/register <user_id>` (ví dụ: `/register john`)
+- User ID phải có ít nhất 3 ký tự
+
+---
+**Ví dụ:**
+- `/login alice`
+- `/register bob123`
+"""
     
-    # Welcome message
-    welcome_msg = """🤖 **Chào mừng bạn đến với Kumiko v2!**
+    await cl.Message(content=welcome_msg, author="Kumiko").send()
+    cl.user_session.set("awaiting_auth", True)
+
+
+async def show_chat_interface(user_id: str):
+    """Display chat interface after successful login"""
+    welcome_msg = f"""#### ✅ Yayy, bạn đã đăng nhập thành công!
 
-Xin chào, mình là Kumiko, trợ lý ảo cute hột me!
+**User ID của bạn là:** `{user_id}`
 
-**Mình có thể giúp bạn những việc sau:**
-- 💬 Trò chuyện và tâm sự với bạn như một người bạn!
-- 🔍 Tìm kiếm và truy xuất thông tin từ Wikipedia và Internet.
+---
 
-Hãy hỏi mình bất kỳ điều gì bạn muốn biết! 😊"""
+Bạn có thể bắt đầu đặt câu hỏi. Mình sẽ trả lời và cung cấp tài liệu tham khảo nếu có!
+
+💡 **Mẹo:** Sau mỗi câu trả lời, bạn có thể nhấn nút "𝄢 Tài liệu tham khảo" để check lại thông tin quan trọng!
+
+---
+Để đăng xuất, gõ `/logout`
+"""
     
-    await cl.Message(
-        content=welcome_msg,
-        author="Kumiko"
-    ).send()
+    await cl.Message(content=welcome_msg, author="Kumiko").send()
+    cl.user_session.set("user_id", user_id)
+    cl.user_session.set("awaiting_auth", False)
+
 
 @cl.on_message
 async def main(message: cl.Message):
-    """Handle incoming messages - only use /chat endpoint"""
-    user_query = message.content.strip()
+    """Handle incoming messages"""
+    user_id = cl.user_session.get("user_id")
+    awaiting_auth = cl.user_session.get("awaiting_auth", True)
     
-    # Validate input
-    if not user_query:
+    content = message.content.strip()
+    
+    # Handle authentication commands
+    if content.startswith("/"):
+        await handle_command(content)
+        return
+    
+    # Check if user is logged in
+    if not user_id or awaiting_auth:
         await cl.Message(
-            content="❗ Vui lòng nhập câu hỏi của bạn.",
+            content="⚠️ Vui lòng đăng nhập trước khi chat.\n\nGõ `/login <user_id>` hoặc `/register <user_id>`",
             author="Kumiko"
         ).send()
         return
     
-    # Show processing indicator
-    async with cl.Step(name="processing", type="run") as step:
-        step.output = "Đang xử lý câu hỏi của bạn..."
-        
-        try:
-            # Send query to /chat endpoint only
-            response = await api_client.chat(user_query)
-            
-            if response.get("success"):
-                agent_response = response.get("response", "Không có phản hồi từ agent")
-                message_count = response.get("message_count", 0)
-                
-                # Update step status
-                step.output = f"✅ Đã xử lý xong (tin nhắn #{message_count})"
-                
-                # Send agent response
-                await cl.Message(
-                    content=agent_response,
-                    author="Kumiko"
-                ).send()
-                
-            else:
-                step.output = "❌ Lỗi xử lý"
-                await cl.Message(
-                    content="❌ Xin lỗi, tôi không thể trả lời câu hỏi này. Vui lòng thử lại với câu hỏi khác.",
-                    author="Kumiko"
-                ).send()
-                
-        except Exception as e:
-            step.output = f"❌ Lỗi: {str(e)}"
-            error_msg = f"❌ **Lỗi khi xử lý:** {str(e)}\n\nVui lòng thử lại hoặc kiểm tra kết nối API."
-            await cl.Message(
-                content=error_msg,
-                author="Kumiko"
-            ).send()
+    # Process chat message
+    await process_chat_message(content, user_id)
 
-@cl.on_chat_end
-async def end():
-    """Handle chat end - automatically call /clear-all when user clicks 'New Chat'"""
-    try:
-        print("🧹 Chat ended - calling /clear-all endpoint...")
+
+async def handle_command(command: str):
+    """Handle special commands"""
+    parts = command.split(maxsplit=1)
+    cmd = parts[0].lower()
+    
+    # Login command
+    if cmd == "/login":
+        if len(parts) < 2:
+            await cl.Message(content="❌ Sử dụng: `/login <user_id>`", author="Kumiko").send()
+            return
+        
+        user_id = parts[1].strip()
+        
+        if len(user_id) < 3:
+            await cl.Message(content="❌ User ID phải có ít nhất 3 ký tự!", author="Kumiko").send()
+            return
         
-        # Call /clear-all endpoint to reset everything
-        result = await api_client.clear_all()
+        # Show loading with spinner
+        async with cl.Step(name="🔍 Đang kiểm tra tài khoản...", type="tool"):
+            # Check if user exists
+            exists = await asyncio.to_thread(check_user_exists, user_id)
         
-        if result.get("success"):
-            print("✅ Successfully cleared all memory and cache")
-            # Optional: Send confirmation message (user might not see it due to chat ending)
+        if exists:
+            await show_chat_interface(user_id)
+        else:
             await cl.Message(
-                content="🧹 **Đã dọn dẹp session hoàn tất** - Sẵn sàng cho cuộc trò chuyện mới!",
+                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"
             ).send()
+    
+    # Register command
+    elif cmd == "/register":
+        if len(parts) < 2:
+            await cl.Message(content="❌ Sử dụng: `/register <user_id>`", author="Kumiko").send()
+            return
+        
+        user_id = parts[1].strip()
+        
+        if len(user_id) < 3:
+            await cl.Message(content="❌ User ID phải có ít nhất 3 ký tự!", author="Kumiko").send()
+            return
+        
+        # Show loading with spinner
+        async with cl.Step(name="📝 Đang đăng ký tài khoản...", type="tool"):
+            # Register user
+            result = await asyncio.to_thread(register_user, user_id)
+        
+        if result["success"]:
+            await cl.Message(content=f"✅ {result['message']}", author="Kumiko").send()
+            await show_chat_interface(user_id)
         else:
-            print("⚠️ Warning: clear-all returned unsuccessful status")
-            
-    except Exception as e:
-        print(f"❌ Error during chat end cleanup: {e}")
-        # Still try to send error message
-        await cl.Message(
-            content=f"⚠️ **Cảnh báo:** Có lỗi khi dọn dẹp session: {str(e)}",
-            author="Kumiko"
-        ).send()
+            await cl.Message(content=f"❌ {result['message']}", author="Kumiko").send()
+    
+    # Logout command
+    elif cmd == "/logout":
+        cl.user_session.set("user_id", None)
+        cl.user_session.set("awaiting_auth", True)
+        await cl.Message(content="👋 Đã đăng xuất thành công!", author="Kumiko").send()
+        await show_login_screen()
+    
+    # Help command
+    elif cmd == "/help":
+        help_msg = """
+#### 📖 Hướng dẫn sử dụng
+
+#### Lệnh đăng nhập/đăng ký:
+- `/login <user_id>` - Đăng nhập với user_id có sẵn
+- `/register <user_id>` - Đăng ký user_id mới
+- `/logout` - Đăng xuất
+
+#### Lệnh khác:
+- `/help` - Hiển thị hướng dẫn này
+
+#### Chat:
+Sau khi đăng nhập, bạn có thể chat bình thường. Hệ thống sẽ trả lời và cung cấp tài liệu tham khảo.
+"""
+        await cl.Message(content=help_msg, author="Kumiko").send()
+    
+    else:
+        await cl.Message(content=f"❌ Lệnh không hợp lệ: `{cmd}`\n\nGõ `/help` để xem danh sách lệnh.", author="Kumiko").send()
+
+
+async def process_chat_message(query: str, user_id: str):
+    """Process chat message and display response"""
+    # Show thinking message with spinner
+    async with cl.Step(name="🤔 Đang suy nghĩ...", type="llm"):
+        # Send request to API
+        result = await asyncio.to_thread(send_chat_message, query, user_id)
+    
+    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()
+
+
+@cl.action_callback("show_references")
+async def on_show_references(action: cl.Action):
+    """Handle reference button click"""
+    # Get references from action payload
+    references = action.payload.get("references", "Không có tài liệu tham khảo")
+    
+    # Send references as a new message
+    await cl.Message(
+        content=references,
+        author="Kumiko"
+    ).send()
+    
+    # Optional: Remove the action button after clicking
+    await action.remove()
+
+
+@cl.on_chat_end
+async def end():
+    """Handle chat end"""
+    print("Chat session ended")
+
+# ============================================
+# CONFIGURATION
+# ============================================
+
+@cl.set_starters
+async def set_starters():
+    """Set starter prompts for quick access"""
+    return [
+        cl.Starter(
+            label="Đăng nhập",
+            message="/login your_user_id",
+            icon="🔑",
+        ),
+        cl.Starter(
+            label="Đăng ký",
+            message="/register your_user_id",
+            icon="✍🏻",
+        ),
+        cl.Starter(
+            label="Trợ giúp",
+            message="/help",
+        ),
+    ]
 
-@cl.on_stop
-async def stop():
-    """Handle application stop - cleanup resources"""
-    try:
-        await api_client.close()
-        print("🔌 API client connection closed")
-    except Exception as e:
-        print(f"Error closing API client: {e}")
 
-if __name__ == "__main__":
-    print("🚀 Starting Chainlit UI for RAG Agent...")
-    print("📡 Connecting to API at:", API_BASE_URL)
-    print("🌐 UI will be available at: http://localhost:8000")
-    cl.run()

chainlit.md

@@ -1,13 +1,76 @@
-# Chào mừng đến với trợ lý ảo Kumiko v2.0 📯🎺🎷
-Tìm hiểu thêm về thiết lập nhân vật của Kumiko: https://hibike-euphonium.fandom.com/wiki/Kumiko_Oumae
+## 🎼 Chào mừng đến với trợ lý ảo Kumiko v2.0 📯🎺🎷
+
+Xin chào! Mình là **Kumiko**, trợ lý ảo được lấy cảm hứng từ nhân vật Kumiko Oumae trong series Hibike! Euphonium.
+
+---
+
+### 🎭 Về Kumiko
+
+Tìm hiểu thêm về thiết lập nhân vật của Kumiko: **[Kumiko Oumae - Hibike! Euphonium Wiki](https://hibike-euphonium.fandom.com/wiki/Kumiko_Oumae)**
+
+---
+
+### ✨ Khả năng của Kumiko
 
 Kumiko hiện cho phép bạn:
-- Tìm kiếm nhân vật, sự vật, sự kiện, khái niệm trên Wikipedia.
-- Tra cứu thông tin, tin tức, thời tiết,... trên Internet.
 
-Sử dụng nút New Chat ở góc trên bên trái để reset phiên trò chuyện.
-Có thể chuyển giữa Light Mode/Dark Mode bằng button góc trên bên phải.
+* 🔍 **Tìm kiếm nhân vật, sự vật, sự kiện, khái niệm** trên Wikipedia
+* 🌐 **Tra cứu thông tin, tin tức, thời tiết,...** trên Internet
+* 💬 **Trò chuyện tự nhiên** và trả lời các câu hỏi của bạn
+* 📚 **Cung cấp tài liệu tham khảo** cho mỗi câu trả lời
+
+---
+
+### 📝 Hướng dẫn sử dụng
+
+#### 1️⃣ Đăng nhập/Đăng ký
+
+Trước khi bắt đầu trò chuyện, bạn cần đăng nhập hoặc đăng ký:
+
+- **Đăng nhập**: `/login <user_id>` (ví dụ: `/login alice`)
+- **Đăng ký**: `/register <user_id>` (ví dụ: `/register bob123`)
+- **User ID** phải có ít nhất 3 ký tự
+
+#### 2️⃣ Bắt đầu trò chuyện
+
+Sau khi đăng nhập thành công, bạn có thể:
+
+- Đặt bất kỳ câu hỏi nào cho Kumiko
+- Nhấn nút **"𝄢 Tài liệu tham khảo"** sau mỗi câu trả lời để xem nguồn thông tin
+- Sử dụng các lệnh đặc biệt bên dưới
+
+#### 3️⃣ Các lệnh hữu ích
+
+- `/help` - Hiển thị hướng dẫn
+- `/logout` - Đăng xuất khỏi phiên hiện tại
+
+---
+
+### 🎯 Mẹo sử dụng
+
+- 🔄 **New Chat**: Sử dụng nút **"New Chat"** ở góc trên bên trái để reset phiên trò chuyện
+- 🌓 **Theme**: Có thể chuyển giữa **Light Mode/Dark Mode** bằng button góc trên bên phải
+- 📚 **References**: Luôn kiểm tra tài liệu tham khảo để xác minh thông tin
+- 💡 **Câu hỏi cụ thể**: Đặt câu hỏi càng rõ ràng, câu trả lời càng chính xác
+
+---
+
+### ⚠️ Lưu ý quan trọng
+
+- Mỗi phiên trò chuyện được liên kết với **User ID** của bạn
+- Thông tin chat được lưu trữ để cải thiện trải nghiệm sử dụng
+- Hãy luôn kiểm tra lại thông tin quan trọng từ nguồn 
+
+---
+
+### 🚀 Sẵn sàng bắt đầu?
+
+Hãy đăng nhập hoặc đăng ký để bắt đầu trò chuyện với Kumiko nhé! 🎺
+
+**Ví dụ:**
+- `/login alice`
+- `/register john_doe`
+
+---
 
-Lưu ý, hiện Kumiko chưa có tính năng ghi nhớ memory theo session ID. 
-Điều này có nghĩa là thông tin cuộc trò chuyện có thể bị rò rỉ cho bên thứ ba.
-Vui lòng không chia sẻ thông tin nhạy cảm. Hãy ứng xử lịch sự.
+💙 *Chúc bạn có trải nghiệm tuyệt vời cùng Kumiko!*