Upload chainlit_hf.py

chainlit_hf.py

@@ -0,0 +1,1100 @@
+import re
+import time
+import chainlit as cl
+import pandas as pd
+import httpx
+import asyncio
+from typing import Dict, List, Any, Optional, Callable
+from dataclasses import dataclass, field
+import os
+import uuid
+from datetime import datetime, timedelta
+
+
+API_BASE_URL = os.getenv("API_BASE_URL")
+
+
+@dataclass
+class ConversationState:
+    """Data class to hold conversation state"""
+    session_id: Optional[str] = None
+    specs_advantages: Dict[str, Any] = field(default_factory=dict)
+    solution_packages: List[str] = field(default_factory=list)
+    raw_documents: Optional[Dict[str, Any]] = None
+    outputs: Optional[Dict[str, Any]] = None
+    selected_model: str = "Gemini 2.0 Flash"
+    product_model_search: bool = False
+    method: str = "dense"  # "dense", "sparse", "hybrid"
+    is_enhance_query: bool = False  # New field for query enhancement toggle
+    enhanced_image_retrieval: bool = False  # New field for enhanced image retrieval toggle
+    # New fields for delayed cleanup - now using asyncio
+    pending_cleanup: bool = False
+    cleanup_task: Optional[asyncio.Task] = None
+    last_activity: datetime = field(default_factory=datetime.now)
+    
+    def reset(self):
+        """Reset state to initial values"""
+        self.session_id = None
+        self.specs_advantages = {}
+        self.solution_packages = []
+        self.raw_documents = None
+        self.outputs = None
+        self.selected_model = "Gemini 2.0 Flash"
+        self.product_model_search = False
+        self.method = "dense"
+        self.is_enhance_query = False
+        self.enhanced_image_retrieval = False
+        # Reset cleanup fields but don't touch tasks
+        self.pending_cleanup = False
+        self.last_activity = datetime.now()
+
+    def cancel_cleanup_task(self):
+        """Cancel pending cleanup task if exists"""
+        if self.cleanup_task and not self.cleanup_task.done():
+            self.cleanup_task.cancel()
+            self.cleanup_task = None
+            print(f"🚫 Cancelled cleanup task for session: {self.session_id}")
+
+
+class StateManager:
+    """Manages conversation state operations with per-session isolation and delayed cleanup"""
+    
+    # CLASS-LEVEL session storage for isolation between different browser sessions
+    _session_states: Dict[str, ConversationState] = {}
+    _lock = asyncio.Lock()  # Async lock for consistency
+    
+    @staticmethod
+    async def get_or_create_session_state(session_id: str) -> ConversationState:
+        """Get existing session state or create new one"""
+        async with StateManager._lock:
+            if session_id not in StateManager._session_states:
+                state = ConversationState()
+                state.session_id = session_id
+                StateManager._session_states[session_id] = state
+                print(f"🆕 Created new session state for: {session_id}")
+            else:
+                state = StateManager._session_states[session_id]
+                print(f"🔄 Retrieved existing session state for: {session_id}")
+                
+                # CRITICAL: If session was pending cleanup, cancel it because user is active again
+                if state.pending_cleanup:
+                    state.cancel_cleanup_task()
+                    state.pending_cleanup = False
+                    print(f"♻️ User activity detected! Cancelled pending cleanup for: {session_id}")
+            
+            # Update activity timestamp
+            state.last_activity = datetime.now()
+            return state
+    
+    @staticmethod
+    async def schedule_delayed_cleanup(session_id: str, delay_seconds: int = 3600):
+        """Schedule delayed cleanup for a session using asyncio (default 1 hour for disconnect tolerance)"""
+        async with StateManager._lock:
+            if session_id not in StateManager._session_states:
+                print(f"⚠️ Cannot schedule cleanup for non-existent session: {session_id}")
+                return
+                
+            state = StateManager._session_states[session_id]
+            
+            # Cancel existing task if any
+            state.cancel_cleanup_task()
+            
+            # Mark as pending cleanup
+            state.pending_cleanup = True
+            
+            # Schedule new cleanup using asyncio
+            async def delayed_cleanup():
+                try:
+                    await asyncio.sleep(delay_seconds)
+                    print(f"⏰ Executing delayed cleanup for session: {session_id}")
+                    await StateManager._perform_actual_cleanup(session_id)
+                except asyncio.CancelledError:
+                    print(f"🚫 Cleanup task cancelled for session: {session_id}")
+                    raise
+                except Exception as e:
+                    print(f"❌ Error in delayed cleanup for {session_id}: {e}")
+            
+            state.cleanup_task = asyncio.create_task(delayed_cleanup())
+            
+            print(f"⏱️ Scheduled cleanup in {delay_seconds}s for session: {session_id} (likely disconnect)")
+    
+    @staticmethod
+    async def _perform_actual_cleanup(session_id: str):
+        """Perform the actual cleanup after delay"""
+        async with StateManager._lock:
+            if session_id not in StateManager._session_states:
+                print(f"⚠️ Session already cleaned or doesn't exist: {session_id}")
+                return
+                
+            state = StateManager._session_states[session_id]
+            
+            # Double-check if session is still pending cleanup (user might have sent message)
+            if not state.pending_cleanup:
+                print(f"🚫 Cleanup cancelled - user activity detected for: {session_id}")
+                return
+            
+            # Perform API cleanup using httpx
+            try:
+                if API_BASE_URL:
+                    payload = {
+                        "reset_cache": True, 
+                        "reset_model": False, 
+                        "session_id": session_id
+                    }
+                    async with httpx.AsyncClient(timeout=30.0) as client:
+                        response = await client.post(f"{API_BASE_URL}/clear_memory", json=payload)
+                        print(f"Clear memory response for {session_id}: {response.status_code}")
+            except Exception as e:
+                print(f"Warning: clear_memory failed for {session_id}: {e}")
+            
+            # Remove from memory
+            del StateManager._session_states[session_id]
+            print(f"🗑️ Successfully cleaned up session: {session_id}")
+    
+    @staticmethod
+    async def cleanup_session_immediate(session_id: str):
+        """Immediate cleanup (for testing or forced cleanup)"""
+        async with StateManager._lock:
+            if session_id in StateManager._session_states:
+                state = StateManager._session_states[session_id]
+                state.cancel_cleanup_task()
+                await StateManager._perform_actual_cleanup(session_id)
+    
+    @staticmethod
+    async def clear_chat_state(state: ConversationState):
+        """Clear all conversation history and reset state via API (but keep session alive)"""
+        if state.session_id is not None and API_BASE_URL:
+            try:
+                payload = {
+                    "reset_cache": True, 
+                    "reset_model": False, 
+                    "session_id": state.session_id
+                }
+                async with httpx.AsyncClient(timeout=30.0) as client:
+                    response = await client.post(f"{API_BASE_URL}/clear_memory", json=payload)
+                    print(f"Clear memory response: {response.status_code}")
+            except Exception as e:
+                print(f"Warning: clear_memory failed: {e}")
+        
+        # Reset state but keep session_id and don't trigger cleanup
+        session_id = state.session_id
+        state.reset()
+        state.session_id = session_id
+    
+    @staticmethod
+    async def change_model(state: ConversationState, model_name: str):
+        """Change the selected model"""
+        state.selected_model = model_name
+        state.last_activity = datetime.now()
+    
+    @staticmethod
+    async def toggle_product_model_search(state: ConversationState):
+        """Toggle product model search mode"""
+        state.product_model_search = not state.product_model_search
+        state.last_activity = datetime.now()
+    
+    @staticmethod
+    async def toggle_enhance_query(state: ConversationState):
+        """Toggle query enhancement mode"""
+        state.is_enhance_query = not state.is_enhance_query
+        state.last_activity = datetime.now()
+
+    @staticmethod
+    async def toggle_enhanced_image_retrieval(state: ConversationState):
+        """Toggle enhanced image retrieval mode"""
+        state.enhanced_image_retrieval = not state.enhanced_image_retrieval
+        state.last_activity = datetime.now()
+    
+    @staticmethod
+    async def cycle_search_method(state: ConversationState):
+        """Cycle search method: dense -> sparse -> hybrid -> dense"""
+        if state.method == "dense":
+            state.method = "sparse"
+        elif state.method == "sparse":
+            state.method = "hybrid"
+        else:
+            state.method = "dense"
+        state.last_activity = datetime.now()
+    
+    @staticmethod
+    async def get_session_status() -> Dict[str, Dict[str, Any]]:
+        """Get status of all sessions (for debugging)"""
+        async with StateManager._lock:
+            status = {}
+            for session_id, state in StateManager._session_states.items():
+                status[session_id] = {
+                    "pending_cleanup": state.pending_cleanup,
+                    "has_task": state.cleanup_task is not None and not state.cleanup_task.done(),
+                    "last_activity": state.last_activity.isoformat(),
+                    "selected_model": state.selected_model,
+                    "product_model_search": state.product_model_search,
+                    "method": state.method,
+                    "is_enhance_query": state.is_enhance_query,
+                    "enhanced_image_retrieval": state.enhanced_image_retrieval
+                }
+            return status
+
+
+class ChatService:
+    """Handles chat-related operations with async HTTP calls"""
+    
+    @staticmethod
+    async def respond_to_chat(
+        state: ConversationState, 
+        message: str, 
+        image_path: Optional[str] = None
+    ) -> str:
+        """Handle chat responses with image support using async HTTP"""
+        print(f"🔄 === DEBUG STATE ===\nChat request with model: {state.selected_model}, Product Model Search: {state.product_model_search}, Method: {state.method}, Session ID: {state.session_id}")
+        
+        # Update activity timestamp - this is KEY to prevent cleanup during active use
+        state.last_activity = datetime.now()
+                
+        if not API_BASE_URL:
+            return "Error: API_BASE_URL not configured"
+        
+        if not state.session_id:
+            return "Error: Session ID not initialized"
+        
+        # Call API using httpx for async HTTP
+        try:
+            async with httpx.AsyncClient(timeout=600.0) as client:
+                if image_path:
+                    # For image uploads, use form-data format as expected by API
+                    with open(image_path, 'rb') as f:
+                        files = {"image": f.read()}
+                    
+                    data = {
+                        "message": message,
+                        "product_model_search": str(state.product_model_search).lower(),
+                        "method": state.method,
+                        "session_id": state.session_id,
+                        "llm_model": state.selected_model,
+                        "debug": "Normal",
+                        "is_enhance_query": str(state.is_enhance_query).lower(),
+                        "enhanced_image_retrieval": str(state.enhanced_image_retrieval).lower()
+                    }
+                    
+                    # Use multipart form data for image upload
+                    files_dict = {"image": ("image.jpg", files["image"], "image/jpeg")}
+                    resp = await client.post(
+                        f"{API_BASE_URL}/chat_with_image", 
+                        files=files_dict, 
+                        data=data
+                    )
+                else:
+                    # For text messages, use form-data format as expected by API
+                    data = {
+                        "message": message,
+                        "session_id": state.session_id,
+                        "debug": "Normal",
+                        "product_model_search": str(state.product_model_search).lower(),
+                        "method": state.method,
+                        "llm_model": state.selected_model,
+                        "is_enhance_query": str(state.is_enhance_query).lower(),
+                        "enhanced_image_retrieval": str(state.enhanced_image_retrieval).lower()
+                    }
+                    resp = await client.post(
+                        f"{API_BASE_URL}/chat",
+                        data=data  # Form data format
+                    )
+
+            if resp.status_code == 200:
+                j = resp.json()
+                response = j.get("response", "")
+                specs_advantages = j.get("specs_advantages")
+                solution_packages = j.get("solution_packages")
+                raw_documents = j.get("raw_documents")  # This might be None from API
+                outputs = j.get("outputs")
+            else:
+                print(f"API Error: {resp.status_code} - {resp.text}")
+                response = f"Error: API status {resp.status_code}"
+                specs_advantages, solution_packages, raw_documents, outputs = None, None, None, None
+        except Exception as e:
+            print(f"Exception calling API: {e}")
+            response = f"Error calling API: {e}"
+            specs_advantages, solution_packages, raw_documents, outputs = None, None, None, None
+
+
+        # Update state
+        if specs_advantages is not None:
+            state.specs_advantages = specs_advantages
+        if solution_packages is not None:
+            state.solution_packages = solution_packages
+        if raw_documents is not None:
+            state.raw_documents = raw_documents
+        if outputs is not None:
+            state.outputs = outputs
+
+        # Filter products based on query
+        if state.specs_advantages is not None:
+            await ChatService.get_specific_product_from_query(message, state)
+
+        # NEW: Format response with 2-column grid for products
+        formatted_response = ChatService.format_product_grid(response)
+
+        return formatted_response
+
+    @staticmethod
+    def format_product_grid(response_text: str) -> str:
+        """Format product listings into 2-column grid while keeping other content intact"""
+        # Pattern to match: *   **[Name](url)**\n\n    ![alt](img_url)
+        pattern = r'\*\s+\*\*\[(.*?)\]\((.*?)\)\*\*\s*\n\s*!\[(.*?)\]\((.*?)\)'
+        
+        matches = list(re.finditer(pattern, response_text))
+        
+        if not matches:
+            # No product listings found, return original
+            return response_text
+        
+        # Find boundaries of product section
+        first_match_start = matches[0].start()
+        last_match_end = matches[-1].end()
+        
+        # Split into: intro + products + rest
+        intro_text = response_text[:first_match_start].strip()
+        rest_text = response_text[last_match_end:].strip()
+        
+        # Extract all products
+        products = []
+        for match in matches:
+            products.append({
+                'name': match.group(1),
+                'url': match.group(2),
+                'alt': match.group(3),
+                'img': match.group(4)
+            })
+        
+        # Build 2-column markdown table
+        grid_content = "\n\n"
+        
+        for i in range(0, len(products), 2):
+            p1 = products[i]
+            
+            if i + 1 < len(products):
+                p2 = products[i + 1]
+                # Two columns
+                grid_content += f"| **[{p1['name']}]({p1['url']})** | **[{p2['name']}]({p2['url']})** |\n"
+                grid_content += f"|:---:|:---:|\n"
+                grid_content += f"| ![{p1['alt']}]({p1['img']}) | ![{p2['alt']}]({p2['img']}) |\n\n"
+            else:
+                # Single column for last odd product
+                grid_content += f"| **[{p1['name']}]({p1['url']})** |\n"
+                grid_content += f"|:---:|\n"
+                grid_content += f"| ![{p1['alt']}]({p1['img']}) |\n\n"
+        
+        # Reconstruct full response
+        return intro_text + grid_content + rest_text
+
+    @staticmethod
+    async def get_specific_product_from_query(query, state):
+        """Filter specs_advantages based on models found in query"""
+        specs_map = state.specs_advantages or {}
+        product_model_list = []
+        
+        for prod_id, data in specs_map.items():
+            model = data.get("model", None) 
+            if model is not None:
+                product_model_list.append(model)
+        
+        found_models = []
+        for model in product_model_list:
+            pattern = re.escape(model)
+            if re.search(pattern, query, re.IGNORECASE):
+                found_models.append(model)
+        
+        new_specs_advantages = {}
+        if found_models != []: 
+            for prod_id, data in specs_map.items():
+                if data.get("model", None) in found_models:
+                    new_specs_advantages[prod_id] = data
+        
+            state.specs_advantages = new_specs_advantages
+
+
+class DisplayService:
+    """Handles display-related operations with async HTTP calls"""
+    
+    @staticmethod
+    async def show_specs(state: ConversationState) -> str:
+        """Generate specifications table"""
+        specs_map = state.specs_advantages
+        columns = ["Thông số"]
+        raw_data = []
+
+        if not specs_map:
+            return "📄 **Thông số kỹ thuật**\n\nKhông có thông số kỹ thuật nào."
+
+        print(specs_map)
+        for prod_id, data in specs_map.items():
+            spec = data.get("specification")
+            if spec is None or spec == "" or spec == "None":
+                spec = "Hiện tại trong dữ liệu chưa có thông tin về thông số kĩ thuật của sản phẩm này!"
+            model = data.get("model", "")
+            url = data.get("url", "")
+
+            # Handle both products and solution packages
+            if url:
+                # full_name = f"**[{data['name']} {model}]({url})**"
+                full_name = f"**[{data['name']}]({url})**"
+            else:
+                # full_name = f"**{data['name']} {model}**"
+                full_name = f"**{data['name']}**"
+
+            if full_name not in columns:
+                columns.append(full_name)
+
+            if spec:
+                # Check if this is a solution package (contains markdown table)
+                if "### 📦" in spec:
+                    # For solution packages, parse the markdown table properly
+                    lines = spec.split('\n')
+                    in_table = False
+                    headers = []
+
+                    for line in lines:
+                        line = line.strip()
+                        if '|' in line and '---' not in line and line.startswith('|') and line.endswith('|'):
+                            cells = [cell.strip()
+                                     for cell in line.split('|')[1:-1]]
+
+                            if not in_table:
+                                # This is the header row
+                                headers = cells
+                                in_table = True
+                                continue
+
+                            # This is a data row
+                            if len(cells) >= len(headers):
+                                for i, header in enumerate(headers):
+                                    if i < len(cells):
+                                        param_name = header
+                                        param_value = cells[i]
+
+                                        existing_row = None
+                                        for row in raw_data:
+                                            if row["Thông số"] == param_name:
+                                                existing_row = row
+                                                break
+
+                                        if existing_row:
+                                            existing_row[full_name] = param_value
+                                        else:
+                                            new_row = {"Thông số": param_name}
+                                            for col in columns[1:]:
+                                                new_row[col] = ""
+                                            new_row[full_name] = param_value
+                                            raw_data.append(new_row)
+                        elif in_table and (not line or not line.startswith('|')):
+                            in_table = False
+                else:
+                    # For products, parse specification items
+                    items = re.split(r';|\n', spec)
+                    for item in items:
+                        if ":" in item:
+                            key, value = item.split(':', 1)
+                            spec_key = key.strip().capitalize()
+                            if spec_key == "Vậtl iệu":
+                                spec_key = "Vật liệu"
+
+                            if "|" in spec_key:
+                                spec_key = spec_key.strip().replace("|", "").capitalize()
+
+                            existing_row = None
+                            for row in raw_data:
+                                if row["Thông số"] == spec_key:
+                                    existing_row = row
+                                    break
+
+                            if existing_row:
+                                existing_row[full_name] = value.strip() if value else ""
+                            else:
+                                new_row = {"Thông số": spec_key}
+                                for col in columns[1:]:
+                                    new_row[col] = ""
+                                new_row[full_name] = value.strip() if value else ""
+                                raw_data.append(new_row)
+
+        if raw_data:
+            df = pd.DataFrame(raw_data, columns=columns)
+            df = df.fillna("").replace("None", "").replace("nan", "")
+        else:
+            df = pd.DataFrame(
+                [["Không có thông số kỹ thuật", "", ""]], columns=columns)
+
+        markdown_table = df.to_markdown(index=False)
+        return f"📄 **Thông số kỹ thuật**\n\n{markdown_table}"
+    
+    @staticmethod
+    async def show_advantages(state: ConversationState) -> str:
+        """Generate advantages as bullet list instead of table"""
+        specs_map = state.specs_advantages
+
+        if not specs_map:
+            return "💡 **Ưu điểm nổi trội**\n\nKhông có ưu điểm nào."
+
+        content = "💡 **Ưu điểm nổi trội**\n\n"
+        
+        for prod_id, data in specs_map.items():
+            # adv = data.get("advantages", "Hiện tại trong dữ liệu chưa có thông tin về ưu điểm nổi trội của sản phẩm này!")
+            adv = data.get("advantages")
+            if adv is None or adv == "" or adv == "None":
+                adv = "Hiện tại trong dữ liệu chưa có thông tin về ưu điểm nổi trội của sản phẩm này!"
+            model = data.get("model", "")
+            url = data.get("url", "")
+
+            # Handle both products and solution packages
+            if url:
+                full_name = f"**[{data['name']}]({url})**"
+            else:
+                full_name = f"**{data['name']}**"
+
+            content += f"### {full_name}\n"
+            
+            # Split by newlines and create bullet points
+            advantages_list = [line.strip() for line in adv.split('\n') if line.strip()]
+            for advantage in advantages_list:
+                content += f"- {advantage}\n"
+            content += "\n"
+
+        return content
+    
+    @staticmethod
+    async def show_solution_packages(state: ConversationState) -> str:
+        """Show solution packages in a structured format"""
+        packages = state.solution_packages
+
+        if not packages or packages == []:
+            return "📦 **Gói sản phẩm**\n\nKhông có gói sản phẩm nào"
+
+        markdown_table = "\n\n".join(packages)
+        return markdown_table
+
+    @staticmethod
+    async def show_all_products_table(state: ConversationState):
+        """Show all products table using async HTTP"""
+        outputs = state.outputs or {}
+
+        if not outputs:
+            return "Không có dữ liệu sản phẩm"
+        
+        try:
+            # Updated to match API format - send outputs in request body
+            payload = {"outputs": outputs}
+            async with httpx.AsyncClient(timeout=60.0) as client:
+                resp = await client.post(f"{API_BASE_URL}/products_by_category", json=payload)
+
+            if resp.status_code == 200:
+                data = resp.json()
+                return data.get("markdown_table", "Không có dữ liệu sản phẩm")
+            else:
+                print(f"All products API error: {resp.status_code} - {resp.text}")
+                return "Không có dữ liệu sản phẩm"
+        except Exception as e:
+            print(f"Exception in show_all_products_table: {e}")
+            return f"Error: {e}"
+
+
+class UIService:
+    """Handles UI-related operations"""
+    
+    @staticmethod
+    def create_action_buttons(state: ConversationState):
+        """Create persistent action buttons"""
+        search_status = "🔍 Tìm theo mã sản phẩm (Đang tắt)" if not state.product_model_search else "🔍 Tìm theo mã sản phẩm (Đang bật)"
+        method_labels = {
+            "dense": "🔎 Tìm kiếm: Dense",
+            "sparse": "🔎 Tìm kiếm: Sparse (BM25)",
+            "hybrid": "🔎 Tìm kiếm: Hybrid"
+        }
+        method_status = method_labels.get(state.method, "🔎 Tìm kiếm: Dense")
+        enhance_status = "🧠 Tăng cường truy vấn (Đang tắt)" if not state.is_enhance_query else "🧠 Tăng cường truy vấn (Đang bật)"
+        enhanced_retrieval_status = "🖼️ Tìm bằng ảnh nâng cao (Đang tắt)" if not state.enhanced_image_retrieval else "🖼️ Tìm bằng ảnh nâng cao (Đang bật)"
+        
+        return [
+            cl.Action(name="show_specs", value="specs", label="📄 Thông số kỹ thuật", payload={"action": "specs"}),
+            cl.Action(name="show_advantages", value="advantages", label="💡 Ưu điểm nổi trội", payload={"action": "advantages"}),
+            cl.Action(name="show_packages", value="packages", label="📦 Gói sản phẩm", payload={"action": "packages"}),
+            cl.Action(name="show_all_products", value="all_products", label="🛒 Tất cả sản phẩm", payload={"action": "all_products"}),
+            cl.Action(name="toggle_product_search", value="toggle_search", label=search_status, payload={"action": "toggle_search"}),
+            cl.Action(name="change_search_method", value="change_method", label="🔎 Đổi phương thức tìm kiếm", payload={"action": "change_method"}),
+            cl.Action(name="toggle_enhance_query", value="toggle_enhance", label=enhance_status, payload={"action": "toggle_enhance"}),
+            cl.Action(name="toggle_enhanced_image_retrieval", value="toggle_enhanced_retrieval", label=enhanced_retrieval_status, payload={"action": "toggle_enhanced_retrieval"}),
+            cl.Action(name="change_model", value="model", label="🔄 Đổi model", payload={"action": "model"}),
+        ]
+
+    @staticmethod
+    def create_start_buttons(state: ConversationState):
+        """Create start buttons"""
+        search_status = "🔍 Tìm theo mã sản phẩm (Đang tắt)" if not state.product_model_search else "🔍 Tìm theo mã sản phẩm (Đang bật)"
+        method_labels = {
+            "dense": "🔎 Tìm kiếm: Dense",
+            "sparse": "🔎 Tìm kiếm: Sparse (BM25)",
+            "hybrid": "🔎 Tìm kiếm: Hybrid"
+        }
+        method_status = method_labels.get(state.method, "🔎 Tìm kiếm: Dense")
+        enhance_status = "🧠 Tăng cường truy vấn (Đang tắt)" if not state.is_enhance_query else "🧠 Tăng cường truy vấn (Đang bật)"
+        enhanced_retrieval_status = "🖼️ Tìm bằng ảnh nâng cao (Đang tắt)" if not state.enhanced_image_retrieval else "🖼️ Tìm bằng ảnh nâng cao (Đang bật)"
+        
+        return [
+            cl.Action(name="toggle_product_search", value="toggle_search", label=search_status, payload={"action": "toggle_search"}),
+            cl.Action(name="change_search_method", value="change_method", label="🔎 Đổi phương thức tìm kiếm", payload={"action": "change_method"}),
+            cl.Action(name="toggle_enhance_query", value="toggle_enhance", label=enhance_status, payload={"action": "toggle_enhance"}),
+            cl.Action(name="toggle_enhanced_image_retrieval", value="toggle_enhanced_retrieval", label=enhanced_retrieval_status, payload={"action": "toggle_enhanced_retrieval"}),
+            cl.Action(name="change_model", value="model", label="🔄 Đổi model", payload={"action": "model"}),
+        ]
+    
+    @staticmethod
+    async def send_message_with_buttons(content: str, state: ConversationState, actions=None, author="assistant"):
+        """Send message with optional action buttons and author"""
+        if actions is None:
+            actions = UIService.create_action_buttons(state)
+        await cl.Message(
+            content=content, 
+            actions=actions,
+            author=author
+        ).send()
+    
+    @staticmethod
+    async def create_typing_animation():
+        """Create typing animation effect (legacy method - kept for compatibility)"""
+        msg = cl.Message(content="", author="assistant")
+        await msg.send()
+        
+        # Typing animation frames
+        typing_frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+        
+        for i in range(27):  # Show animation for ~2 seconds
+            frame = typing_frames[i % len(typing_frames)]
+            msg.content = f"{frame} Đang suy nghĩ..."
+            await msg.update()
+            await asyncio.sleep(0.25)
+        
+        return msg
+
+
+async def run_typing_animation(msg: cl.Message):
+    """Run typing animation until cancelled"""
+    typing_frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+    frame_index = 0
+    
+    try:
+        while True:  # Run indefinitely until cancelled
+            frame = typing_frames[frame_index % len(typing_frames)]
+            msg.content = f"{frame} Đang suy nghĩ..."
+            await msg.update()
+            await asyncio.sleep(0.25)
+            frame_index += 1
+            
+    except asyncio.CancelledError:
+        # Animation was cancelled, this is expected
+        print("🎬 Animation cancelled - API response received")
+        raise
+
+
+# HELPER FUNCTIONS: Session management with proper async error handling
+async def ensure_session_state() -> Optional[ConversationState]:
+    """Ensure session state exists, create if not"""
+    try:
+        session_id = cl.user_session.get("session_id")
+        
+        if not session_id:
+            print(f"Lỗi: Không lấy được session id ở ensure_session_state")
+            return None
+        
+        return await StateManager.get_or_create_session_state(session_id)
+        
+    except Exception as e:
+        print(f"⚠️ Error ensuring session state: {e}")
+        return None
+
+
+async def get_current_session_state() -> Optional[ConversationState]:
+    """Get current session state using Chainlit's session system"""
+    try:
+        # Use Chainlit's user session to get unique session ID
+        chainlit_session_id = cl.user_session.get("session_id")
+        
+        if chainlit_session_id:
+            return await StateManager.get_or_create_session_state(chainlit_session_id)
+        else:
+            print("⚠️ No Chainlit session ID found")
+            return None
+    except Exception as e:
+        print(f"⚠️ Error getting session state: {e}")
+        return None
+
+
+@cl.on_chat_start
+async def on_chat_start():
+    """Initialize the chat session"""
+    session_id = cl.user_session.get("session_id")
+    if not session_id:
+        session_id = str(uuid.uuid4())
+        cl.user_session.set("session_id", session_id)
+        print(f"🆕 Generated new session_id: {session_id}")
+    else:
+        print(f"🔄 Reusing existing session_id: {session_id}")
+        
+    app_state = await StateManager.get_or_create_session_state(session_id)
+
+    await cl.Message(
+        content=f"🛍️ **RangDong Sales Agent** (Session: {session_id[:8]}...)\n\n"
+                   f"Xin chào! Tôi có thể giúp bạn tìm kiếm và tư vấn sản phẩm RangDong. Hãy thử các câu hỏi mẫu:\n\n"
+                   f"- Tìm sản phẩm bình giữ nhiệt dung tích dưới 2 lít\n"
+                   f"- Tìm sản phẩm ổ cắm thông minh\n"
+                   f"- Tư vấn cho tôi đèn học chống cận cho con gái của tôi học lớp 6",
+        author="assistant"
+    ).send()
+
+    actions = UIService.create_start_buttons(app_state)
+    await cl.Message(
+        content="Sử dụng nút bên dưới để cấu hình:", 
+        actions=actions,
+        author="assistant"
+    ).send()
+
+
+@cl.on_chat_end
+async def on_chat_end():
+    """Handle chat session end with delayed cleanup mechanism using asyncio"""
+    try:
+        session_id = cl.user_session.get("session_id")
+        print(f"📤 on_chat_end triggered for session {session_id}")
+        
+        if session_id:
+            # Schedule delayed cleanup instead of immediate cleanup
+            # Use shorter delay (30s) since this is likely just a temporary disconnect
+            await StateManager.schedule_delayed_cleanup(session_id, delay_seconds=3600)
+            print(f"⏳ Scheduled delayed cleanup for session {session_id} (1h delay for disconnect tolerance)")
+        else:
+            print("⚠️ No session_id found in on_chat_end")
+    except Exception as e:
+        print(f"⚠️ Error during on_chat_end: {e}")
+
+
+# ACTION CALLBACKS - All use ensure_session_state() for better reliability
+@cl.action_callback("show_specs")
+async def on_show_specs(action):
+    """Handle show specifications action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    specs_content = await DisplayService.show_specs(app_state)
+    await UIService.send_message_with_buttons(specs_content, app_state, author="assistant")
+
+
+@cl.action_callback("show_advantages")
+async def on_show_advantages(action):
+    """Handle show advantages action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    adv_content = await DisplayService.show_advantages(app_state)
+    await UIService.send_message_with_buttons(adv_content, app_state, author="assistant")
+
+
+@cl.action_callback("show_packages")
+async def on_show_packages(action):
+    """Handle show packages action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    pkg_content = await DisplayService.show_solution_packages(app_state)
+    await UIService.send_message_with_buttons(pkg_content, app_state, author="assistant")
+
+@cl.action_callback("show_all_products")
+async def on_show_all_products(action):
+    """Handle show all products action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    all_products_content = await DisplayService.show_all_products_table(app_state)
+    await UIService.send_message_with_buttons(all_products_content, app_state, author="assistant")
+
+@cl.action_callback("toggle_product_search")
+async def on_toggle_product_search(action):
+    """Handle toggle product model search action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    await StateManager.toggle_product_model_search(app_state)
+    
+    status_message = (
+        "✅ **Đã bật tìm kiếm theo mã sản phẩm**\n\n"
+        "Khi bạn nhắc đến mã/model cụ thể trong câu hỏi, hệ thống sẽ tìm kiếm chính xác theo mã đó."
+        if app_state.product_model_search 
+        else "✅ **Đã tắt tìm kiếm theo mã sản phẩm**\n\n"
+        "Hệ thống sẽ tìm kiếm sản phẩm theo cách thông thường."
+    )
+    
+    await UIService.send_message_with_buttons(status_message, app_state, author="assistant")
+
+
+@cl.action_callback("toggle_enhanced_image_retrieval")
+async def on_toggle_enhanced_image_retrieval(action):
+    """Handle toggle enhanced image retrieval action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    await StateManager.toggle_enhanced_image_retrieval(app_state)
+    
+    status_message = (
+        "✅ **Đã bật tìm bằng ảnh nâng cao**\n\n"
+        "Hệ thống sẽ sử dụng Gemini để phân tích kỹ hình ảnh và tạo từ khóa tìm kiếm chi tiết."
+        if app_state.enhanced_image_retrieval 
+        else "✅ **Đã tắt tìm bằng ảnh nâng cao**\n\n"
+        "Hệ thống sẽ sử dụng tìm kiếm hình ảnh thông thường (Visual Semantic Search)."
+    )
+    
+    await UIService.send_message_with_buttons(status_message, app_state, author="assistant")
+
+
+@cl.action_callback("toggle_enhance_query")
+async def on_toggle_enhance_query(action):
+    """Handle toggle enhance query action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    await StateManager.toggle_enhance_query(app_state)
+    
+    status_message = (
+        "✅ **Đã bật tăng cường truy vấn**\n\n"
+        "Hệ thống sẽ tự động cải thiện và mở rộng câu hỏi của bạn để tìm kiếm chính xác hơn."
+        if app_state.is_enhance_query 
+        else "✅ **Đã tắt tăng cường truy vấn**\n\n"
+        "Hệ thống sẽ sử dụng câu hỏi gốc của bạn mà không cải thiện."
+    )
+    
+    await UIService.send_message_with_buttons(status_message, app_state, author="assistant")
+
+
+@cl.action_callback("change_search_method")
+async def on_change_search_method(action):
+    """Handle change search method action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    method_actions = [
+        cl.Action(name="select_method_dense", value="dense", label="🔎 Dense (Mặc định)", payload={"method": "dense"}),
+        cl.Action(name="select_method_sparse", value="sparse", label="🔎 Sparse (BM25)", payload={"method": "sparse"}),
+        cl.Action(name="select_method_hybrid", value="hybrid", label="🔎 Hybrid", payload={"method": "hybrid"}),
+        cl.Action(name="back_to_main", value="back", label="🔙 Quay lại", payload={"action": "back"})
+    ]
+    
+    current_method_labels = {
+        "dense": "Dense",
+        "sparse": "Sparse (BM25)",
+        "hybrid": "Hybrid"
+    }
+    current = current_method_labels.get(app_state.method, "Dense")
+    
+    await cl.Message(
+        content=f"**Model hiện tại**: {app_state.selected_model}\n**Tìm kiếm theo mã**: {'Đang bật' if app_state.product_model_search else 'Đang tắt'}\n**Phương thức tìm kiếm**: {current}\n**Tăng cường truy vấn**: {'Đang bật' if app_state.is_enhance_query else 'Đang tắt'}\n\nChọn phương thức tìm kiếm mới:",
+        actions=method_actions,
+        author="assistant"
+    ).send()
+
+
+@cl.action_callback("select_method_dense")
+async def on_select_method_dense(action):
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    app_state.method = "dense"
+    app_state.last_activity = datetime.now()
+    await UIService.send_message_with_buttons("✅ Đã chuyển sang **Dense Search**\n\nHệ thống sẽ sử dụng tìm kiếm semantic vector thông thường.", app_state, author="assistant")
+
+
+@cl.action_callback("select_method_sparse")
+async def on_select_method_sparse(action):
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    app_state.method = "sparse"
+    app_state.last_activity = datetime.now()
+    await UIService.send_message_with_buttons("✅ Đã chuyển sang **Sparse Search (BM25)**\n\nHệ thống sẽ sử dụng tìm kiếm từ khóa BM25.", app_state, author="assistant")
+
+
+@cl.action_callback("select_method_hybrid")
+async def on_select_method_hybrid(action):
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    app_state.method = "hybrid"
+    app_state.last_activity = datetime.now()
+    await UIService.send_message_with_buttons("✅ Đã chuyển sang **Hybrid Search**\n\nHệ thống sẽ kết hợp cả Dense và Sparse vector.", app_state, author="assistant")
+
+
+@cl.action_callback("change_model")
+async def on_change_model(action):
+    """Handle model change action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    models = ["Gemini 2.0 Flash", "Gemini 2.5 Flash Lite", "Gemini 2.0 Flash Lite"]
+    
+    model_actions = [
+        cl.Action(name=f"select_model_{i}", value=model, label=model, payload={"model": model})
+        for i, model in enumerate(models)
+    ]
+    
+    model_actions.append(
+        cl.Action(name="back_to_main", value="back", label="🔙 Quay lại", payload={"action": "back"})
+    )
+    
+    await cl.Message(
+        content=f"**Model hiện tại**: {app_state.selected_model}\n**Tìm kiếm theo mã**: {'Đang bật' if app_state.product_model_search else 'Đang tắt'}\n**Phương thức tìm kiếm**: {app_state.method}\n**Tăng cường truy vấn**: {'Đang bật' if app_state.is_enhance_query else 'Đang tắt'}\n\nChọn model mới:",
+        actions=model_actions,
+        author="assistant"
+    ).send()
+
+
+@cl.action_callback("back_to_main")
+async def on_back_to_main(action):
+    """Handle back to main menu action"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    actions = UIService.create_action_buttons(app_state)
+    await cl.Message(
+        content="📋 **Menu chính**\n\nSử dụng các nút bên dưới để:", 
+        actions=actions,
+        author="assistant"
+    ).send()
+
+
+@cl.action_callback("select_model_0")
+async def on_select_model_0(action):
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    await StateManager.change_model(app_state, "Gemini 2.0 Flash")
+    await UIService.send_message_with_buttons("✅ Đã chuyển sang **Gemini 2.0 Flash**", app_state, author="assistant")
+
+
+@cl.action_callback("select_model_1")
+async def on_select_model_1(action):
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    await StateManager.change_model(app_state, "Gemini 2.5 Flash Lite")
+    await UIService.send_message_with_buttons("✅ Đã chuyển sang **Gemini 2.5 Flash Lite**", app_state, author="assistant")
+
+
+@cl.action_callback("select_model_2")
+async def on_select_model_2(action):
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+        
+    await StateManager.change_model(app_state, "Gemini 2.0 Flash Lite")
+    await UIService.send_message_with_buttons("✅ Đã chuyển sang **Gemini 2.0 Flash Lite**", app_state, author="assistant")
+
+
+# DEBUG ENDPOINTS (optional - for monitoring session status)
+@cl.action_callback("debug_sessions")
+async def on_debug_sessions(action):
+    """Debug action to show session status (can be added to debug builds)"""
+    try:
+        status = await StateManager.get_session_status()
+        debug_content = "🔍 **Debug: Session Status**\n\n"
+        
+        if not status:
+            debug_content += "No active sessions."
+        else:
+            for session_id, info in status.items():
+                debug_content += f"**Session: {session_id[:8]}...**\n"
+                debug_content += f"- Pending cleanup: {info['pending_cleanup']}\n"
+                debug_content += f"- Has task: {info['has_task']}\n"
+                debug_content += f"- Last activity: {info['last_activity']}\n"
+                debug_content += f"- Model: {info['selected_model']}\n"
+                debug_content += f"- Product search: {info['product_model_search']}\n"
+                debug_content += f"- Method: {info.get('method', 'dense')}\n\n"
+        
+        await cl.Message(content=debug_content, author="assistant").send()
+    except Exception as e:
+        await cl.Message(content=f"Debug error: {e}", author="assistant").send()
+
+
+@cl.on_message
+async def main(message: cl.Message):
+    """Main message handler with concurrent animation and API call"""
+    app_state = await ensure_session_state()
+    if app_state is None:
+        await cl.Message(content="Error: Session state not found", author="assistant").send()
+        return
+    
+    # Handle images if present
+    image_path = None
+    if message.elements:
+        for element in message.elements:
+            if isinstance(element, cl.Image):
+                image_path = element.path
+                break
+    
+    # Create initial message for animation
+    msg = cl.Message(content="", author="assistant")
+    await msg.send()
+    
+    # Create concurrent tasks for animation and API call
+    animation_task = asyncio.create_task(run_typing_animation(msg))
+    api_task = asyncio.create_task(ChatService.respond_to_chat(app_state, message.content, image_path))
+    
+    try:
+        # Wait for API response (this will complete first usually)
+        response = await api_task
+        
+        # Cancel animation task since we have the response
+        animation_task.cancel()
+        
+        # Wait a bit for graceful animation cancellation
+        try:
+            await asyncio.wait_for(animation_task, timeout=0.1)
+        except (asyncio.CancelledError, asyncio.TimeoutError):
+            pass
+        
+    except Exception as e:
+        # If API fails, cancel animation and show error
+        animation_task.cancel()
+        try:
+            await asyncio.wait_for(animation_task, timeout=0.1)
+        except (asyncio.CancelledError, asyncio.TimeoutError):
+            pass
+        response = f"Error: {e}"
+    
+    # Update message with final response and buttons
+    msg.content = response
+    msg.actions = UIService.create_action_buttons(app_state)
+    await msg.update()