Update app.py

app.py

@@ -1,6 +1,1141 @@
+from llama_index.llms.ollama import Ollama
+from llama_index.embeddings.huggingface_optimum import OptimumEmbedding
+from llama_index.core import Settings
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.storage.chat_store import SimpleChatStore
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.duckdb import DuckDBVectorStore
+from llama_index.core.llms import ChatMessage, MessageRole
+import uuid
+import os
+import json
+import nest_asyncio
+from datetime import datetime
+import copy
+import ollama
+
 import gradio as gr
+from gradio.themes.utils import colors, fonts, sizes
+from gradio.themes import Base
+from huggingface_hub import whoami
+import re
+
+
+
+
+from llama_index.core.evaluation import FaithfulnessEvaluator
+
+from huggingface_hub import snapshot_download
+import html
+import concurrent.futures
+import time
+
+nest_asyncio.apply()
+
+
+PERSISTENT_DIR = "/data"
+
+FORCE_UPDATE_FLAG = False
+
+
+
+
+VECTOR_STORE_DIR = "./vector_stores"
+EMBED_MODEL_PATH = "./datas/bge_onnx"
+CONFIG_PATH = "config.json"
+
+DEFAULT_LLM = "Jatin19K/unsloth-q5_k_m-mistral-nemo-instruct-2407:latest"
+DEFAULT_VECTOR_STORE = "CFIR"
+
+CONVERSATION_HISTORY_PATH = "./conversation_history"
+
+
+SYSTEM_PROMPT = (
+    "You are a helpful assistant which helps users to understand scientific knowledge "
+    "about biomechanics of injuries to human bodies."
+)
+
+
+# HF required
+EMBED_MODEL_PATH = os.path.join(PERSISTENT_DIR, "bge_onnx")
+VECTOR_STORE_DIR = os.path.join(PERSISTENT_DIR, "vector_stores")
+CONVERSATION_HISTORY_PATH = os.path.join(PERSISTENT_DIR, "conversation_history")
+token = os.getenv("HF_TOKEN")
+dataset_id = os.getenv("DATASET_ID")
+
+def download_data_if_needed():
+    global FORCE_UPDATE_FLAG
+    
+    if not os.path.exists(EMBED_MODEL_PATH) or not os.path.exists(VECTOR_STORE_DIR):
+        FORCE_UPDATE_FLAG = True
+
+    if FORCE_UPDATE_FLAG:
+        snapshot_download(
+                repo_id=dataset_id,
+                repo_type="dataset",
+                token=token,
+                local_dir=PERSISTENT_DIR
+            )
+        print("Data downloaded successfully.")
+    else:
+        print("Data exists.")
+
+download_data_if_needed()
+
+
+
+
+def process_text_with_think_tags(text):
+    # Check if the text contains think tags
+    think_pattern = r'<think>(.*?)</think>'
+    think_matches = re.findall(think_pattern, text, re.DOTALL)
+    
+    if think_matches:
+        # There are think tags present
+        # Extract the content inside think tags
+        think_content = think_matches[0]  # Taking the first think block
+        
+        # Remove the think tags part from the original text
+        remaining_text = re.sub(think_pattern, '', text, flags=re.DOTALL).strip()
+        
+        # Return both parts separately
+        return {
+            'has_two_parts': True,
+            'think_part': think_content,
+            'regular_part': remaining_text
+        }
+    else:
+        # No think tags, just one part
+        return {
+            'has_two_parts': False,
+            'full_text': text
+        }
+    
+    
+
+
+class VectorStoreManager:
+    def __init__(self):
+        self.vector_stores = self.initialize_vector_stores()
+
+
+    def initialize_vector_stores(self):
+        """Scan vector store directory for DuckDB files, supporting nested directories"""
+        vector_stores = {}
+        if os.path.exists(VECTOR_STORE_DIR):
+            # Add default store if it exists
+            cfir_path = os.path.join(VECTOR_STORE_DIR, f"{DEFAULT_VECTOR_STORE}.duckdb")
+            if os.path.exists(cfir_path):
+                vector_stores[DEFAULT_VECTOR_STORE] = {
+                    "path": cfir_path,
+                    "display_name": DEFAULT_VECTOR_STORE,
+                    "data": DuckDBVectorStore.from_local(cfir_path)
+                }
+            
+            # Scan for .duckdb files in root directory and subdirectories
+            for root, dirs, files in os.walk(VECTOR_STORE_DIR):
+                for file in files:
+                    if file.endswith(".duckdb") and file != f"{DEFAULT_VECTOR_STORE}.duckdb":
+                        # Skip the default store since we've already handled it
+                        if root == VECTOR_STORE_DIR and file == f"{DEFAULT_VECTOR_STORE}.duckdb":
+                            continue
+                        
+                        # Get the full path to the file
+                        file_path = os.path.join(root, file)
+                        
+                        # Calculate store_name: combine category and subcategory
+                        rel_path = os.path.relpath(file_path, VECTOR_STORE_DIR)
+                        path_parts = rel_path.split(os.sep)
+                        
+                        if len(path_parts) == 1:
+                            # Files in the root directory
+                            store_name = path_parts[0][:-7]  # Remove .duckdb
+                            display_name = store_name
+                        else:
+                            # Files in subdirectories
+                            category = path_parts[0]
+                            file_name = path_parts[-1][:-7]  # Remove .duckdb
+                            store_name = f"{category}_{file_name}"
+                            display_name = f"{category} - {file_name}"
+                        
+                        vector_stores[store_name] = {
+                            "path": file_path,
+                            "display_name": display_name,
+                            "data": DuckDBVectorStore.from_local(file_path)
+                        }
+        
+        return vector_stores
+
+
+    def get_vector_store_data(self, store_name):
+        """Get the actual vector store data by store name"""
+        return self.vector_stores[store_name]["data"]
+    
+    def get_vector_store_by_display_name(self, display_name):
+        """Find a vector store by its display name"""
+        for name, store_info in self.vector_stores.items():
+            if store_info["display_name"] == display_name:
+                return self.vector_stores[name]["data"]
+        return None
+    
+    def get_all_store_names(self):
+        """Get all vector store names"""
+        return list(self.vector_stores.keys())
+    
+    def get_all_display_names(self):
+        """Get all display names as a list"""
+        return [store_info["display_name"] for store_info in self.vector_stores.values()]
+    
+    def get_display_name(self, store_name):
+        """Get display name for a store name"""
+        return self.vector_stores[store_name]["display_name"]
+    
+    def get_name_display_pairs(self):
+        """Get list of (display_name, store_name) tuples for UI dropdowns"""
+        return [(v["display_name"], k) for k, v in self.vector_stores.items()]
+
+# Create a global instance
+vector_store_manager = VectorStoreManager()
+
+
+
+class DEKCIBChatbot:
+    def __init__(self):
+        self.initialize()
+
+
+    def initialize(self):
+        self.session_manager = SessionManager()
+        self.embed_model = OptimumEmbedding(folder_name=EMBED_MODEL_PATH)
+        Settings.embed_model = self.embed_model
+        self.vector_stores = self.initialize_vector_store()
+
+        self.config = self._load_config()
+        self.llm_options = self._initialize_models()
+
+
+
+    def get_user_data(self, user_id):
+        return user_id
+
+
+
+    
+
+    def _load_config(self):
+        """Load model configuration from JSON file"""
+        try:
+            with open(CONFIG_PATH, 'r') as f:
+                return json.load(f)
+        except Exception as e:
+            print(f"Error loading config: {e}")
+            return {"models": []}
+
+    def _initialize_models(self):
+        """Initialize and verify all models from config"""
+        config_models = self.config.get("models", [])
+        available_models = {}
+
+        # Get currently available Ollama models
+        try:
+            current_models = {m['name']: m['name'] for m in ollama.list()['models']}
+            print(current_models)
+        except Exception as e:
+            print(f"Error fetching current models: {e}")
+            current_models = {}
+
+        # Check each configured model
+        for model_name in config_models:
+            if model_name not in current_models:
+                print(f"Model {model_name} not found locally. Attempting to pull...")
+                try:
+                    ollama.pull(model_name)
+                    available_models[model_name] = model_name
+                    print(f"Successfully pulled model {model_name}")
+                except Exception as e:
+                    print(f"Error pulling model {model_name}: {e}")
+                    continue
+            else:
+                available_models[model_name] = current_models[model_name]
+
+        return available_models
+
+    def get_available_models(self):
+        """Return dictionary of available models"""
+        return self.available_models
+    
+
+    def initialize_vector_store(self):
+        """Scan vector store directory for DuckDB files, supporting nested directories"""
+        vector_stores = {}
+        if os.path.exists(VECTOR_STORE_DIR):
+            # Add default store if it exists
+            cfir_path = os.path.join(VECTOR_STORE_DIR, f"{DEFAULT_VECTOR_STORE}.duckdb")
+            if os.path.exists(cfir_path):
+                vector_stores[DEFAULT_VECTOR_STORE] = {
+                    "path": cfir_path,
+                    "display_name": DEFAULT_VECTOR_STORE,
+                    "data": DuckDBVectorStore.from_local(cfir_path)
+                }
+            
+            # Scan for .duckdb files in root directory and subdirectories
+            for root, dirs, files in os.walk(VECTOR_STORE_DIR):
+                for file in files:
+                    if file.endswith(".duckdb") and file != f"{DEFAULT_VECTOR_STORE}.duckdb":
+                        # Skip the default store since we've already handled it
+                        if root == VECTOR_STORE_DIR and file == f"{DEFAULT_VECTOR_STORE}.duckdb":
+                            continue
+                        
+                        # Get the full path to the file
+                        file_path = os.path.join(root, file)
+                        
+                        # Calculate store_name: combine category and subcategory
+                        rel_path = os.path.relpath(file_path, VECTOR_STORE_DIR)
+                        path_parts = rel_path.split(os.sep)
+                        
+                        if len(path_parts) == 1:
+                            # Files in the root directory
+                            store_name = path_parts[0][:-7]  # Remove .duckdb
+                            display_name = store_name
+                        else:
+                            # Files in subdirectories
+                            category = path_parts[0]
+                            file_name = path_parts[-1][:-7]  # Remove .duckdb
+                            store_name = f"{category}_{file_name}"
+                            display_name = f"{category} - {file_name}"
+                        
+                        vector_stores[store_name] = {
+                            "path": file_path,
+                            "display_name": display_name,
+                            "data": DuckDBVectorStore.from_local(file_path)
+                        }
+        
+        return vector_stores
+    
+    
+    def get_vector_store(self, vector_store_name):
+        return self.vector_stores[vector_store_name]["data"]
+
+
+class DeKCIBChatEngine:
+    """
+    Manages the core components needed for chat functionality with RAG.
+    Handles LLM, vector store, memory, chat store, and indexes.
+    """
+    
+    def __init__(self, user_id=None, llm_name=None, vector_store_name=None):
+        """Initialize the chat engine with all necessary components"""
+        self.user_id = user_id
+        self.llm = None
+        self.llm_name = llm_name
+        self.vector_store = None
+        self.vector_store_name = vector_store_name
+        self.storage_context = None
+        self.index = None
+        self.chat_store = None
+        self.memory = None
+        self.chat_engine = None
+        self.rebuild_chat_engine_flag = True
+
+        
+        # Conversation metadata management
+        self.convs_metadata = {}
+        self.current_conv_id = None
+
+        if user_id:
+            self.initialize_chat_store()
+            self.initialize_convs_metadata()
+        
+        # Set initial components if provided
+        if llm_name:
+            self.set_llm(llm_name)
+        
+        if vector_store_name:
+            self.set_vector_store(vector_store_name)
+
+
+
+    def initialize_convs_metadata(self):
+        print(f"Initializing convs metadata for user {self.user_id}")
+        self.convs_metadata_file_path = os.path.join(CONVERSATION_HISTORY_PATH, self.user_id, f"{self.user_id}_metadata.json")
+        self.sorted_conversation_list = []
+        self.get_convs_metadata()
+
+
+
+    def get_convs_metadata(self):
+        if os.path.exists(self.convs_metadata_file_path):
+            with open(self.convs_metadata_file_path, "r") as f:
+                self.convs_metadata = json.load(f)
+            self.sorted_conversation_list = self.get_sorted_conversation_list()
+
+
+
+
+    def set_current_conv_id(self, input_value, type="index"):
+
+        if len(self.sorted_conversation_list) == 0:
+            self.current_conv_id = None
+            self.rebuild_chat_engine_flag = True
+            return
+
+        if type == "index" and self.current_conv_id != self.sorted_conversation_list[input_value]:
+            self.current_conv_id = self.sorted_conversation_list[input_value]
+            self.rebuild_chat_engine_flag = True
+        elif type == "id" and self.current_conv_id != input_value:
+            self.current_conv_id = input_value
+            self.rebuild_chat_engine_flag = True
+
+
+
+    def get_sorted_conversation_list(self):
+        """
+        Returns a list of conversation IDs sorted by update time,
+        with the most recently updated conversations first.
+        """
+        # Create a list of (conv_id, updated_at) tuples
+        conv_with_timestamps = []
+        
+        for conv_id, metadata in self.convs_metadata.items():
+            # Use updated_at timestamp for sorting
+            if "updated_at" in metadata:
+                # Convert the ISO timestamp string to datetime object for comparison
+                update_time = datetime.fromisoformat(metadata["updated_at"])
+                conv_with_timestamps.append((conv_id, update_time))
+        
+        # Sort by timestamp (descending order - newest first)
+        sorted_convs = sorted(conv_with_timestamps, key=lambda x: x[1], reverse=True)
+        
+        # Return just the conversation IDs in the sorted order
+        return [conv_id for conv_id, _ in sorted_convs]
+
+
+    def get_sorted_conversation_list_for_ui(self):
+        new_list = []
+        for item in self.sorted_conversation_list:
+            new_list.append([self.convs_metadata[item]["title"]])
+        return new_list
+
+
+    def update_convs_metadata(self, conv_id, title=None, create_flag=False):
+        current_time = datetime.now().isoformat()
+        if title is not None:
+            self.convs_metadata[conv_id].update({"title":title})
+        self.convs_metadata[conv_id].update({"updated_at":current_time, "llm_name": self.llm_name, "vector_store_name": self.vector_store_name})
+
+        self.sorted_conversation_list = self.get_sorted_conversation_list()
+
+
+    
+    def set_llm(self, llm_name):
+
+        self.llm = Ollama(
+            model=llm_name,
+            request_timeout=120,
+            temperature=0.3
+        )
+        self.set_rebuild_chat_engine_flag(True)
+
+        return self.llm
+    
+    def set_vector_store(self, vector_store_name):
+
+        self.vector_store = vector_store_manager.get_vector_store_by_display_name(vector_store_name)
+            
+        if self.vector_store:
+            self.initialize_index()        
+            self.set_rebuild_chat_engine_flag(True)
+
+        return self.vector_store
+    
+    def initialize_index(self):
+        """Initialize the index using the current vector store"""
+        if not self.vector_store:
+            raise ValueError("Vector store must be set before initializing index")
+            
+        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
+        self.index = VectorStoreIndex.from_vector_store(
+            vector_store=self.vector_store,
+            storage_context=self.storage_context
+        )
+        return self.index
+    
+    def initialize_chat_store(self):
+        """Initialize the chat store for the user"""
+        print(f"Initializing chat store for user {self.user_id}")
+            
+        chat_store_file_path = os.path.join(CONVERSATION_HISTORY_PATH, self.user_id, f"{self.user_id}.json")
+        
+        # Ensure directory exists
+        os.makedirs(os.path.dirname(chat_store_file_path), exist_ok=True)
+        
+        # Create or load chat store
+        if not os.path.exists(chat_store_file_path):
+            self.chat_store = SimpleChatStore()
+            self.chat_store.persist(persist_path=chat_store_file_path)
+        else:
+            self.chat_store = SimpleChatStore.from_persist_path(chat_store_file_path)
+            
+        self.chat_store_file_path = chat_store_file_path
+
+        return self.chat_store
+    
+    
+    def initialize_memory(self, conversation_id=None):
+        """Initialize or reinitialize memory with specified conversation ID"""
+        if not self.chat_store:
+            raise ValueError("Chat store must be initialized before memory")
+        
+
+        print(f"Initializing memory for conversation {conversation_id}")
+        print(self.chat_store)
+
+        self.memory = ChatMemoryBuffer.from_defaults(
+            token_limit=3000,
+            chat_store=self.chat_store,
+            chat_store_key=conversation_id
+        )
+        return self.memory
+    
+    def build_chat_engine(self, conversation_id=None):
+        """Build the chat engine with all components"""
+        if not all([self.llm, self.index, self.chat_store]):
+            raise ValueError("LLM, index, and chat store must be set before building chat engine")
+            
+        # Initialize or update memory with conversation ID
+        # if conversation_id and self.current_conv_id != conversation_id:
+        self.initialize_memory(conversation_id)
+        self.current_conv_id = conversation_id
+        
+        # Default system prompt if none provided
+
+        # Create the chat engine
+        self.chat_engine = self.index.as_chat_engine(
+            chat_mode="context",
+            llm=self.llm,
+            memory=self.memory,
+            system_prompt=SYSTEM_PROMPT
+        )
+
+        self.set_rebuild_chat_engine_flag(False)
+        return self.chat_engine
+    
+    def save_chat_history(self):
+        """Save chat history to file"""
+        if self.chat_store and hasattr(self, 'chat_store_file_path'):
+            self.chat_store.persist(persist_path=self.chat_store_file_path)
+            
+    def add_message(self, conversation_id, message):
+        """Add a message to the chat history"""
+        if self.chat_store:
+            self.chat_store.add_message(conversation_id, message)
+            
+    def get_chat_history(self, conversation_id):
+        """Get chat history for a specific conversation"""
+        if conversation_id is None:
+            return []
+        if self.chat_store:
+            return self.chat_store.to_dict()["store"][conversation_id]
+        return []
+    
+
+    def set_rebuild_chat_engine_flag(self, flag):
+        self.rebuild_chat_engine_flag = flag
+
+    def chat(self, message, conversation_id=None):
+
+
+        create_flag = False
+        if conversation_id is None:
+            conversation_id = self.create_conversation(message=message)
+            create_flag = True
+            print(f"Created new conversation {conversation_id}")
+            self.set_rebuild_chat_engine_flag(True)
+        elif self.current_conv_id != conversation_id:
+            self.set_rebuild_chat_engine_flag(True)
+
+
+
+        if self.rebuild_chat_engine_flag:
+            self.chat_engine = self.build_chat_engine(conversation_id)
+            self.rebuild_chat_engine_flag = False
+
+
+        print("user message")
+        # user_msg = ChatMessage(role=MessageRole.USER, content=message)
+        # self.add_message(conversation_id, user_msg)
+    
+
+        print("L597")
+        print(message)
+        # Get response
+        response = self.chat_engine.chat(message)
+
+        answer = response.response
+
+        print(answer)
+        print(type(answer))
+        print("assistant message")
+        # assistant_msg = ChatMessage(role=MessageRole.ASSISTANT, content=answer)
+        # self.add_message(conversation_id, assistant_msg)
+
+
+        self.update_convs_metadata(conversation_id, create_flag=create_flag)
+        print("update_convs_metadata")
+        self.save_metadata()
+        print("save_metadata")
+        self.save_chat_history()
+        print("save_chat_history")
+        
+        return response
+    
+    def create_conversation(self, message=None):
+        """
+        Create a new conversation with metadata
+        Args:
+            title: Optional title for the conversation
+            message: First message to use for generating a title
+        Returns:
+            conversation_id: ID of the new conversation
+        """
+        # Generate a new unique conversation ID
+        conv_id = str(uuid.uuid4())
+        
+        # Set as current conversation
+        self.current_conv_id = conv_id
+        
+        # Generate title from message if not provided
+
+
+        title = message[:50] + ("..." if len(message) > 50 else "")
+        
+        # Create timestamp
+        current_time = datetime.now().isoformat()
+        
+        # Store metadata with resource information
+        self.convs_metadata[conv_id] = {
+            "title": title,
+            "created_at": current_time,
+            "updated_at": current_time,
+            "llm": self.llm_name,
+            "vector_store": self.vector_store_name,
+            "message_count": 0
+        }
+        
+        # Initialize chat engine with the new conversation ID
+        # self.chat_engine = self.build_chat_engine(conv_id)
+        
+        return conv_id
+    
+    def update_conversation_metadata(self, conv_id, title=None, increment_message_count=True):
+        """
+        Update conversation metadata
+        Args:
+            conv_id: Conversation ID to update
+            title: Optional new title
+            increment_message_count: Whether to increment message count
+        """
+        if conv_id not in self.convs_metadata:
+            return
+        
+        # Update timestamp
+        self.convs_metadata[conv_id]["updated_at"] = datetime.now().isoformat()
+        
+        # Update title if provided
+        if title:
+            self.convs_metadata[conv_id]["title"] = title
+        
+        # Increment message count if requested
+        if increment_message_count:
+            self.convs_metadata[conv_id]["message_count"] = self.convs_metadata[conv_id].get("message_count", 0) + 1
+    
+    def get_sorted_conversations(self):
+        """
+        Returns a list of conversation IDs sorted by update time,
+        with the most recently updated conversations first.
+        """
+        # Create a list of (conv_id, updated_at) tuples
+        conv_with_timestamps = []
+        
+        for conv_id, metadata in self.convs_metadata.items():
+            # Use updated_at timestamp for sorting
+            if "updated_at" in metadata:
+                # Convert the ISO timestamp string to datetime object for comparison
+                update_time = datetime.fromisoformat(metadata["updated_at"])
+                conv_with_timestamps.append((conv_id, update_time))
+        
+        # Sort by timestamp (descending order - newest first)
+        sorted_convs = sorted(conv_with_timestamps, key=lambda x: x[1], reverse=True)
+        
+        # Return just the conversation IDs in the sorted order
+        return [conv_id for conv_id, _ in sorted_convs]
+    
+    def get_conversation_info(self, conv_id):
+        """Get conversation metadata"""
+        return self.convs_metadata.get(conv_id, {})
+    
+    def switch_conversation(self, conv_id):
+        """
+        Switch to an existing conversation
+        Args:
+            conv_id: Conversation ID to switch to
+        Returns:
+            True if successful, False otherwise
+        """
+        if conv_id not in self.convs_metadata:
+            return False
+        
+        # Set as current conversation
+        self.current_conv_id = conv_id
+        
+        # Get the conversation's LLM and vector store
+        metadata = self.convs_metadata[conv_id]
+        
+        # Switch to the conversation's resources if they're different
+        if metadata.get("llm") and metadata["llm"] != self.llm_name:
+            self.set_llm(metadata["llm"])
+            
+        if metadata.get("vector_store") and metadata["vector_store"] != self.vector_store_name:
+            self.set_vector_store(metadata["vector_store"])
+        
+        # Rebuild chat engine with this conversation ID
+        self.build_chat_engine(conv_id)
+        
+        return True
+    
+    def save_metadata(self):
+        """Save conversation metadata to file"""
+        if hasattr(self, 'chat_store_file_path') and self.user_id:
+            metadata_path = os.path.join(CONVERSATION_HISTORY_PATH, self.user_id, f"{self.user_id}_metadata.json")
+            os.makedirs(os.path.dirname(metadata_path), exist_ok=True)
+            with open(metadata_path, 'w') as f:
+                json.dump(self.convs_metadata, f)
+    
+    def load_metadata(self):
+        """Load conversation metadata from file"""
+        if self.user_id:
+            metadata_path = os.path.join(CONVERSATION_HISTORY_PATH, self.user_id, f"{self.user_id}_metadata.json")
+            if os.path.exists(metadata_path):
+                try:
+                    with open(metadata_path, 'r') as f:
+                        self.convs_metadata = json.load(f)
+                except Exception as e:
+                    print(f"Error loading metadata: {e}")
+
+    def get_or_build_chat_engine(self, conversation_id=None, llm_name=None, vector_store_name=None):
+        """
+        Check if the chat engine needs to be rebuilt based on changes to LLM, vector store or conversation ID.
+        Only rebuilds the chat engine if necessary to avoid performance overhead.
+        
+        Args:
+            conversation_id: The conversation ID to use
+            llm_name: The LLM model name to use
+            vector_store_name: The vector store name to use
+            system_prompt: Custom system prompt (optional)
+            
+        Returns:
+            The existing or newly built chat engine
+        """
+        rebuild_needed = False
+        
+        # Check if conversation ID changed
+        if conversation_id is not None and self.conv_id != conversation_id:
+            print(f"Building chat engine: Conversation ID changed from {self.conv_id} to {conversation_id}")
+            self.conv_id = conversation_id
+            rebuild_needed = True
+        
+        # Check if LLM changed
+        if llm_name is not None and self.llm_name != llm_name:
+            print(f"Building chat engine: LLM changed from {self.llm_name} to {llm_name}")
+            self.set_llm(llm_name)
+            rebuild_needed = True
+        
+        # Check if vector store changed
+        if vector_store_name is not None and self.vector_store_name != vector_store_name:
+            print(f"Building chat engine: Vector store changed from {self.vector_store_name} to {vector_store_name}")
+            self.set_vector_store(vector_store_name)
+            rebuild_needed = True
+        
+        # Rebuild only if needed
+        if rebuild_needed:
+            return self.build_chat_engine(conversation_id)
+        else:
+            print("Using existing chat engine: No changes detected")
+            return self.chat_engine
+
+
+
+
+
+class SessionManager:
+    def __init__(self):
+        self.sessions = {}
+
+    def create_session(self, user_id=None):
+        if user_id is None:
+            return None
+        
+        print(f"Creating session for user {user_id}")
+        if user_id not in self.sessions:
+            self.sessions[user_id] = DeKCIBChatEngine(user_id, llm_name=DEFAULT_LLM, vector_store_name=DEFAULT_VECTOR_STORE)
+            print(f"Session created for user {user_id}")
+        return self.sessions[user_id]
+
+
+
+
+
+class ChatbotUI:
+    """UI handler for the chatbot application"""
+    
+    def __init__(self, dekcib_chatbot):
+        """Initialize with a chat engine"""
+        self.dekcib_chatbot = dekcib_chatbot
+        self.init_attr()
+
+
+    def init_attr(self):
+        self.llm_options = self.dekcib_chatbot.llm_options
+        self.vector_stores = self.dekcib_chatbot.vector_stores
+        # self.vector_stores_options = [(v["display_name"], k) for k, v in self.dekcib_chatbot.vector_stores.items()]
+
+
+        # self.init_conversations_history()
+        
+
+    # def init_conversations_history(self):
+    #     chat_session = self.dekcib_chatbot.session_manager.sessions[USER_NAME]
+    #     self.init_convs_list = chat_session.get_sorted_conversation_list_for_ui()
+    #     if len(self.init_convs_list) > 0:
+    #         self.init_chat_history = chat_session.get_chat_history(chat_session.sorted_conversation_list[0])
+    #         self.init_convs_index = 0
+    #     else:
+    #         self.init_chat_history = []
+    #         self.init_convs_index = None
+
+
+
+
+
+
+
+    def create_ui(self):
+        with gr.Blocks(title="De-KCIB(Deep Knowledge Center for Injury Biomechanics)") as demo:
+
+            user_id = gr.State(None)
+
+            with gr.Row():
+                with gr.Column(scale=6):
+                    gr.Markdown("<img src='/gradio_api/file/logo.png' alt='Innovision Logo' height='150' width='390'>")
+                with gr.Column(scale=1):
+                    login_btn = gr.LoginButton()
+
+            with gr.Row():
+                gr.Markdown("# De-KCIB(Deep Knowledge Center for Injury Biomechanics)")
+            
+
+                
+            # Move model selection to the top row
+            with gr.Row():
+                with gr.Column(scale=3):
+                    llm_dropdown = gr.Dropdown(
+                        label="Select Language Model",
+                        choices=list(self.llm_options.values()),
+                        value=next(iter(self.llm_options.values()), None)
+                    )
+                with gr.Column(scale=3):
+                    vector_dropdown = gr.Dropdown(
+                        label="Injury Biomechanics Knowledge Base",
+                        choices=[(v["display_name"], k) for k, v in self.vector_stores.items()],
+                        value=next(iter(self.vector_stores.keys()), None)
+
+                    )
+                with gr.Column(scale=1):
+                    status_indicator = gr.HTML(
+                        value='<div style="text-align:center; padding:8px; border-radius:4px; background-color:#f0f0f0; margin-top:18px;">✓ Ready</div>',
+                        elem_id="status_indicator"
+                    )
+            
+            # Main content with sidebar and chat area
+            with gr.Row():
+                # Left sidebar for conversation history
+                with gr.Column(scale=1, elem_classes="sidebar"):
+                    new_chat_btn = gr.Button("New Chat", size="sm")
+                    
+                    # Hidden textbox for conversation data
+                    conversation_data = gr.Textbox(visible=False)
+
+
+                    # Dataset for conversation history
+                    conversation_history = gr.Dataset(
+                        components=[conversation_data],
+                        label="Conversation History",
+                        type="index",
+                        layout="table"
+                    )
+
+                
+                # Main chat area
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(
+                        height=500, 
+                        render_markdown=True, 
+                        show_copy_button=True,
+                        type="messages",
+                    )
+            
+                    with gr.Row():
+                        msg = gr.Textbox(label="Query", scale=5)
+                        clear_btn = gr.Button("Clear Session", scale=1)
+            
+
+
+
+            def get_auth_id(oauth_token: gr.OAuthToken | None) -> str:
+                if oauth_token is None:
+                    return None                
+                id = whoami(oauth_token.token)['id']
+                return id
+    
+
+
+
+
+            def add_msg(msg, history):
+
+
+                history.append({"role": "user", "content": msg})
+                return history
+
+
+
+
+            def chat_with_dekcib(history, user_id, conv_idx):
+                
+
+                msg = history[-1]["content"]
+
+                user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
+                # user_engine.che
+                history.append({"role": "assistant", "content": ""})
+
+                print("conv_idx")
+                print(conv_idx)
+                
+
+                conv_id = None
+                if conv_idx is not None:
+                    conv_id = user_engine.sorted_conversation_list[conv_idx]
+
+                if len(history) == 1:
+                    conv_id = None
+
+
+
+                response = user_engine.chat(msg, conv_id)
+                answer = response.response
+                for character in answer:
+                    history[-1]["content"] += character
+                    yield history
+
+            def clear_msg():
+                print("clear_msg")
+                return ""
+            
+
+            def update_conversation_history(user_id):
+                user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
+                ui_list = user_engine.get_sorted_conversation_list_for_ui()
+
+                if len(ui_list) > 0:
+                    idx = 0
+                else:
+                    idx = None
+
+                return gr.update(samples=ui_list, value=idx)            
+
+            
+
+            msg.submit(
+                add_msg,
+                [msg, chatbot],
+                [chatbot]
+            ).then(
+                clear_msg,
+                None,
+                [msg]
+            ).then(
+                chat_with_dekcib,
+                [chatbot, user_id, conversation_history],
+                [chatbot]
+            ).then(
+                update_conversation_history,
+                [user_id],
+                [conversation_history]
+            )
+
+
+            def click_to_select_conversation(conversation_history, user_id):
+                user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
+                user_engine.set_current_conv_id(conversation_history, type="index")
+
+                chat_history = user_engine.get_chat_history(user_engine.current_conv_id)
+
+                return gr.update(value=conversation_history), chat_history
+
+            
+            conversation_history.click(
+                click_to_select_conversation,
+                [conversation_history, user_id],
+                [conversation_history, chatbot]
+            )
+
+
+            # msg.submit(
+            #     chat_with_dekcib,
+            #     [msg, chatbot, user_id_dropdown],
+            #     [chatbot]
+            # )
+
+
+            
+            # msg.submit(
+            #     clear_msg,
+            #     None,
+            #     [msg]
+            # ).then(
+            #     chat_with_dekcib,
+            #     [msg, chatbot, user_id_dropdown],
+            #     [chatbot]
+            # )
+
+
+            # clear_btn.click(
+            #     clear_session,
+            #     [session_state],
+            #     [chatbot, session_state],
+            #     queue=False
+            # )
+
+
+            def create_session(user_id):
+                if user_id is None:
+                    return
+
+
+                self.dekcib_chatbot.session_manager.create_session(user_id)
+                user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
+
+                llm_name = user_engine.llm_name
+                vector_store_name = user_engine.vector_store_name
+                # chat_store = user_engine.chat_store
+
+                # convs = user_engine.convs
+                # history = user_engine.history
+
+                sorted_conversation_list = user_engine.get_sorted_conversation_list_for_ui()
+                print("sorted_conversation_list")
+                print(sorted_conversation_list)
+                # sorted_conversation_list = [
+                #     ["I think therefore I am."],
+                #     ["The unexamined life is not worth living."],
+                #     ["Test Item"]
+                # ]
+
+                if len(sorted_conversation_list) > 0:
+                    index = 0
+                else:
+                    index = None
+                update_conversation_history = gr.update(samples=sorted_conversation_list, value=index)
+
+                user_engine.set_current_conv_id(0, type="index")
+                chat_history = user_engine.get_chat_history(user_engine.current_conv_id)
+
+
+                yield llm_name, vector_store_name, update_conversation_history, chat_history
+
+
+            demo.load(
+                get_auth_id, 
+                inputs=None, 
+                outputs=[user_id]
+            ).then(
+                create_session,
+                [user_id],
+                [llm_dropdown, vector_dropdown, conversation_history, chatbot]
+            )
+
+
+            def update_llm(user_id, llm_name):
+                if user_id is None:
+                    return
+
+                user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
+                user_engine.set_llm(llm_name)
+
+
+            llm_dropdown.change(
+                update_llm,
+                [user_id, llm_dropdown],
+                None
+            )
+
+
+            def update_vector_store(user_id, vector_store_name):
+                if user_id is None:
+                    return
+
+                user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
+                user_engine.set_vector_store(vector_store_name)
+
+
+
+            vector_dropdown.change(
+                update_vector_store,
+                [user_id, vector_dropdown],
+                None
+            )
+            
+            
+            
+            # Create new conversation button should only clear the chat area, but not create a new conversation yet
+            def prepare_new_chat():
+                print("prepare_new_chat")
+
+                return [], gr.update(value=None)
+            
+            def print_dataset(value):
+                print("value")
+                print(value)
+            
+            # Create new conversation
+            new_chat_btn.click(
+                prepare_new_chat,
+                None,
+                [chatbot, conversation_history],
+            ).then(
+                print_dataset,
+                conversation_history,
+                None
+            )
+
+        return demo
+
+
+
+# Deployment settings
+if __name__ == "__main__":
+    # Check chat store health
+    # store_health_ok = check_chat_store_health()
+    # if not store_health_ok:
+    #     print("WARNING: Chat store health check failed! Some functionality may not work correctly.")
+    
+    # # Run warm-up to pre-initialize resources
+    # warm_up_resources()
+    
+    dekcib_chatbot = DEKCIBChatbot()
+    ui = ChatbotUI(dekcib_chatbot)
+    demo = ui.create_ui()
+    demo.queue(max_size=10, default_concurrency_limit=3)
+    demo.launch(allowed_paths=["logo.png"])
 
-with gr.Blocks() as demo:
-    button = gr.LoginButton("Sign in")
-    gr.load("models/meta-llama/Llama-3.2-3B-Instruct", provider="sambanova", accept_token=button)
-demo.launch()