Upload 5 files

Dockerfile

@@ -0,0 +1,78 @@
+# Base image with CUDA 12.1 and Ubuntu 22.04
+FROM nvidia/cuda:12.1.1-base-ubuntu22.04
+
+# Install Python 3.10 and essential dependencies
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    python3.10 \
+    python3.10-dev \
+    python3.10-distutils \
+    curl \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Make Python 3.10 the default
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1
+
+# Install pip for Python 3.10
+RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10
+
+
+# Install Ollama with GPU layers
+ENV OLLAMA_GPU_LAYERS=100
+RUN curl -fsSL https://ollama.com/install.sh | sh
+
+# Set up application directory
+WORKDIR /app
+
+
+
+
+# Configure environment variables (FROM YOUR ORIGINAL SETUP)
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH \
+    # VECTOR_STORE_DIR=/app/vector_stores \
+    # EMBED_MODEL_PATH=/app/datas/bge_onnx \
+    PYTHONUNBUFFERED=1 \
+    GRADIO_SERVER_NAME="0.0.0.0" \
+    HF_HOME=/data/.huggingface \
+    HF_HUB_DISABLE_PROGRESS_BARS=1 \
+    OLLAMA_MODELS=/data/.ollama/models \
+    SYSTEM=spaces
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+USER root
+RUN mkdir -p /data && chown user:user /data
+
+# Switch to the "user" user
+WORKDIR /app
+USER user
+COPY --chown=user requirements.txt /app/
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+
+COPY --chown=user . /app
+
+
+# COPY --chown=user . $HOME/app
+# RUN mkdir -p /data
+# RUN chmod 777 /data
+
+
+# Verify CUDA and Python versions
+# RUN python3 -c "import torch; print(f'PyTorch CUDA available: {torch.cuda.is_available()}')" && \
+#     python3 --version
+
+# Expose ports for Ollama and Gradio
+EXPOSE 11434 7860
+
+
+# Copy and set permissions for start script
+# COPY start.sh /app/start.sh
+RUN chmod +x /app/start.sh
+
+# Start services using the startup script
+CMD ["/app/start.sh"]

app.py

@@ -0,0 +1,1332 @@
+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 gradio.events import EditData
+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 = "ComFit"
+
+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
+            comfit_path = os.path.join(VECTOR_STORE_DIR, f"{DEFAULT_VECTOR_STORE}.duckdb")
+            if os.path.exists(comfit_path):
+                vector_stores[DEFAULT_VECTOR_STORE] = {
+                    "path": comfit_path,
+                    "display_name": DEFAULT_VECTOR_STORE,
+                    "data": DuckDBVectorStore.from_local(comfit_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 ComFitChatbot:
+    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
+            comfit_path = os.path.join(VECTOR_STORE_DIR, f"{DEFAULT_VECTOR_STORE}.duckdb")
+            if os.path.exists(comfit_path):
+                vector_stores[DEFAULT_VECTOR_STORE] = {
+                    "path": comfit_path,
+                    "display_name": DEFAULT_VECTOR_STORE,
+                    "data": DuckDBVectorStore.from_local(comfit_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 comfitChatEngine:
+    """
+    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)
+
+
+        self.llm_name = llm_name        
+        if self.current_conv_id:
+            self.convs_metadata[self.current_conv_id].update({"llm_name":self.llm_name})
+
+        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)
+
+            self.vector_store_name = vector_store_name
+
+
+            if self.current_conv_id:
+                self.convs_metadata[self.current_conv_id].update({"vector_store_name":self.vector_store_name})        
+
+        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}")
+
+
+        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 get_chat_history_for_ui(self, conversation_id):
+        """Get chat history for a specific conversation"""
+        if conversation_id is None:
+            return []
+        if self.chat_store:
+            conv_data = self.chat_store.to_dict()["store"][conversation_id]
+
+            conv_data_for_ui = []
+            for item in conv_data:
+                if item["role"] == "user":
+                    conv_data_for_ui.append(item)
+                else:
+
+                    content = item["content"]
+
+
+                    time_str = None
+                    if "time" in item["additional_kwargs"]:
+                        elapsed_time = item["additional_kwargs"]["time"]
+                        time_str = f"\n\n[Total time: {elapsed_time:.2f}s]"
+
+                    processed_answer_dict = process_text_with_think_tags(content)
+
+                    if processed_answer_dict["has_two_parts"]:
+                        think_content = processed_answer_dict["think_part"]
+                        conv_data_for_ui.append({"role": "assistant", "content": think_content, "metadata":{"title":"Thinking...", "status":"done"}})
+                        remaining_text = processed_answer_dict["regular_part"]
+                        if time_str:
+                            remaining_text += time_str
+                        conv_data_for_ui.append({"role": "assistant", "content": remaining_text})
+                    else:
+                        item_copy = copy.deepcopy(item)
+                        if time_str:
+                            item_copy["content"] += time_str
+                        conv_data_for_ui.append(item_copy)
+            return conv_data_for_ui
+
+        return []
+    
+
+    def set_rebuild_chat_engine_flag(self, flag):
+        self.rebuild_chat_engine_flag = flag
+
+    def chat(self, message, conversation_id=None):
+
+        start_time = time.time()
+        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
+
+
+        # Get response
+        response = self.chat_engine.chat(message)
+
+        # answer = response.response
+
+        elapsed_time = time.time() - start_time
+
+
+
+        answer_dict = self.chat_store.get_messages(conversation_id)[-1].dict()
+        answer_dict['additional_kwargs'].update({"time":elapsed_time})
+        new_msg = ChatMessage.model_validate(answer_dict)
+        self.chat_store.delete_message(conversation_id, -1)
+        self.chat_store.add_message(conversation_id, new_msg)
+
+        self.update_convs_metadata(conversation_id, create_flag=create_flag)
+
+        self.save_metadata()
+
+        self.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 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 edit_message(self, index, conversation_id):
+        if conversation_id is not None:
+
+            msg_list = self.chat_store.get_messages(conversation_id)
+            new_msg_list = msg_list[:index]
+
+            self.chat_store.set_messages(conversation_id, new_msg_list)
+
+
+            self.save_metadata()
+            self.save_chat_history()
+
+
+    def retry_message(self, conversation_id):
+        if conversation_id is not None:
+            self.undo_message(conversation_id)
+            self.save_metadata()
+            self.save_chat_history()
+
+
+
+    def undo_message(self, conversation_id):
+        if conversation_id is not None:
+            msg_list = self.chat_store.get_messages(conversation_id)
+
+
+            if msg_list[-1].role == MessageRole.ASSISTANT and len(msg_list) > 0:
+                self.chat_store.delete_last_message(conversation_id)
+            if msg_list[-1].role == MessageRole.USER and len(msg_list) > 0:
+                self.chat_store.delete_last_message(conversation_id)
+
+            
+            
+
+            self.update_convs_metadata(conversation_id)
+            self.save_metadata()
+            self.save_chat_history()
+
+
+    def delete_conversation(self, conversation_id):
+        if conversation_id is not None:
+            self.chat_store.delete_messages(conversation_id)
+            self.convs_metadata.pop(conversation_id)
+            self.save_metadata()
+            self.save_chat_history()
+            self.sorted_conversation_list = self.get_sorted_conversation_list()
+
+
+
+
+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] = comfitChatEngine(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, comfit_chatbot):
+        """Initialize with a chat engine"""
+        self.comfit_chatbot = comfit_chatbot
+        self.init_attr()
+
+
+    def init_attr(self):
+        self.llm_options = self.comfit_chatbot.llm_options
+        self.vector_stores = self.comfit_chatbot.vector_stores
+        # self.vector_stores_options = [(v["display_name"], k) for k, v in self.comfit_chatbot.vector_stores.items()]
+
+
+        # self.init_conversations_history()
+        
+
+    # def init_conversations_history(self):
+    #     chat_session = self.comfit_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="Comfort and Fit Copilot (ComFit Copilot)") 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("# Comfort and Fit Copilot (ComFit Copilot)")
+            
+
+                
+            # 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"]) for k, v in self.vector_stores.items()],
+                        value=next(iter(self.vector_stores.keys()), None)
+
+                    )
+
+            # 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="Ask me anything", placeholder="Log in to start chatting", interactive=False)
+
+
+            # 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 get_auth_id(oauth_token: gr.OAuthToken | None) -> str | None:
+                
+                print(oauth_token)
+                if oauth_token is None:
+                    return None
+                
+                try:
+                    user_info = whoami(oauth_token.token)
+                    print(user_info)
+                    return user_info.get('id')
+                except Exception as e:
+                    print(f"Authentication failed: {e}")
+                    return None
+
+
+
+            def add_msg(msg, history):
+
+
+                history.append({"role": "user", "content": msg})
+                return history
+
+
+
+            def chat_with_comfit(history, user_id, conv_idx):
+                
+                start_time = time.time()
+
+                msg = history[-1]["content"]
+
+                user_engine = self.comfit_chatbot.session_manager.sessions[user_id]
+                # user_engine.che
+                
+                # conv_id = None
+                if conv_idx is not None:
+                    conv_id = user_engine.sorted_conversation_list[conv_idx]
+                else:
+                    conv_id = None
+
+                # if len(history) == 1 and conv_idx is None:
+                #     conv_id = None
+
+
+                response = user_engine.chat(msg, conv_id)
+                answer = response.response
+
+                processed_answer_dict = process_text_with_think_tags(answer)
+
+                                
+                if processed_answer_dict["has_two_parts"]:
+                    think_content = processed_answer_dict["think_part"]
+                    remaining_text = processed_answer_dict["regular_part"]
+
+
+                    # thick_msg = gr.ChatMessage(role="assistant", content="", metadata={"title":"Thinking..."})
+                    history.append({"role": "assistant", "content": "", "metadata":{"title":"Thinking...", "status":"pending"}})
+                    # history.append(thick_msg)
+                    for character in think_content:
+                        history[-1]["content"] += character
+                        yield history
+
+
+                    elapsed_time = time.time() - start_time
+                    history[-1]["metadata"]["title"] = f"Thinking... [Thinking time: {elapsed_time:.2f}s]"
+                    history[-1]["metadata"]["status"] = "done"
+                    yield history
+
+                    # Start response time measurement
+
+                    history.append({"role": "assistant", "content": ""})
+                    for character in remaining_text:
+                        history[-1]["content"] += character
+                        yield history
+
+                    elapsed_time = time.time() - start_time
+                    history[-1]["content"] += f"\n\n[Total time: {elapsed_time:.2f}s]"
+                    yield history
+
+                       
+                else:
+                    full_text = processed_answer_dict["full_text"]
+                    history.append({"role": "assistant", "content": ""})
+                    for character in full_text:
+                        history[-1]["content"] += character
+                        yield history
+                    
+                    elapsed_time = time.time() - start_time
+                    history[-1]["content"] += f"\n\n[Total time: {elapsed_time:.2f}s]"
+                    yield history
+
+
+            def clear_msg():
+                return ""
+            
+
+            def update_conversation_history(user_id):
+                user_engine = self.comfit_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_comfit,
+                [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.comfit_chatbot.session_manager.sessions[user_id]
+                user_engine.set_current_conv_id(conversation_history, type="index")
+
+
+                chat_history = user_engine.get_chat_history_for_ui(user_engine.current_conv_id)
+
+                llm_name = user_engine.convs_metadata[user_engine.current_conv_id]["llm_name"]
+                vector_store_name = user_engine.convs_metadata[user_engine.current_conv_id]["vector_store_name"]
+
+                return gr.update(value=conversation_history), chat_history,  gr.update(value=llm_name), gr.update(value=vector_store_name)
+
+            
+
+            
+
+
+
+            conversation_history.click(
+                click_to_select_conversation,
+                [conversation_history, user_id],
+                [conversation_history, chatbot, llm_dropdown, vector_dropdown]
+            )
+
+
+            # msg.submit(
+            #     chat_with_comfit,
+            #     [msg, chatbot, user_id_dropdown],
+            #     [chatbot]
+            # )
+
+
+            
+            # msg.submit(
+            #     clear_msg,
+            #     None,
+            #     [msg]
+            # ).then(
+            #     chat_with_comfit,
+            #     [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.comfit_chatbot.session_manager.create_session(user_id)
+                user_engine = self.comfit_chatbot.session_manager.sessions[user_id]
+
+
+
+                sorted_conversation_list = user_engine.get_sorted_conversation_list_for_ui()
+
+
+                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_for_ui(user_engine.current_conv_id)
+
+                if len(chat_history) > 0:
+                    llm_name = user_engine.convs_metadata[user_engine.current_conv_id]["llm_name"]
+                    vector_store_name = user_engine.convs_metadata[user_engine.current_conv_id]["vector_store_name"]
+                else:
+                    llm_name = user_engine.llm_name
+                    vector_store_name = user_engine.vector_store_name
+
+                yield llm_name, vector_store_name, update_conversation_history, chat_history
+
+
+            def activate_chat(user_id):
+                if user_id is None:
+                    return gr.update(placeholder="Log in to start chatting", interactive=False)
+                return gr.update(placeholder="",interactive=True)
+
+
+            demo.load(
+                get_auth_id, 
+                inputs=None, 
+                outputs=[user_id]
+            ).then(
+                create_session,
+                [user_id],
+                [llm_dropdown, vector_dropdown, conversation_history, chatbot]
+            ).success(
+                activate_chat,
+                [user_id],
+                [msg]
+            )
+
+
+            def update_llm(user_id, llm_name):
+                if user_id is None:
+                    return
+
+                user_engine = self.comfit_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.comfit_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
+            )
+
+
+            def edit_chat(user_id, history, edit_data: EditData):
+                user_engine = self.comfit_chatbot.session_manager.sessions[user_id]
+                idx = edit_data.index
+                
+                # Count how many user messages appear up to this index in the UI history
+                user_message_count = 0
+                for i in range(idx + 1):
+                    if history[i]["role"] == "user":
+                        user_message_count += 1
+                
+                # In backend storage, user messages are at positions 0, 2, 4, 6...
+                # So the backend index is (user_message_count - 1) * 2
+                backend_idx = (user_message_count - 1) * 2
+                
+                user_engine.edit_message(backend_idx, user_engine.current_conv_id)
+                history = history[: idx+1]
+                return history
+            
+
+            chatbot.edit(
+                edit_chat,
+                [user_id, chatbot],
+                [chatbot]
+                ).success(
+                    chat_with_comfit,
+                    [chatbot, user_id, conversation_history],
+                    [chatbot]
+                ).success(
+                    update_conversation_history,
+                    [user_id],
+                    [conversation_history]
+                )
+
+
+
+            def retry_chat(user_id, history):
+                user_engine = self.comfit_chatbot.session_manager.sessions[user_id]
+                user_engine.retry_message(user_engine.current_conv_id)
+
+
+                while history[-1]["role"] == "assistant":
+                    history.pop()
+                    yield history
+
+ 
+                return history
+
+
+
+
+
+            chatbot.retry(
+                retry_chat,
+                [user_id, chatbot],
+                [chatbot]
+                ).then(
+                    chat_with_comfit,
+                    [chatbot, user_id, conversation_history],
+                    [chatbot]
+                ).then(
+                    update_conversation_history,
+                    [user_id],
+                    [conversation_history]
+                )
+
+
+            def undo_chat(user_id):
+                user_engine = self.comfit_chatbot.session_manager.sessions[user_id]
+                user_engine.undo_message(user_engine.current_conv_id)
+
+                chat_history = user_engine.get_chat_history_for_ui(user_engine.current_conv_id)
+
+                return chat_history
+
+            chatbot.undo(
+                undo_chat,
+                [user_id],
+                [chatbot]
+                )
+            
+
+
+            def clear_conversation(user_id):
+                user_engine = self.comfit_chatbot.session_manager.sessions[user_id]
+                user_engine.delete_conversation(user_engine.current_conv_id)
+
+                sorted_conversation_list = user_engine.get_sorted_conversation_list_for_ui()
+
+                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(index, type="index")
+                chat_history = user_engine.get_chat_history_for_ui(user_engine.current_conv_id)
+
+                yield update_conversation_history, chat_history
+
+
+
+            chatbot.clear(
+                clear_conversation,
+                [user_id],
+                [conversation_history, chatbot]
+                )
+
+            # 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)
+            
+            # 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()
+    
+    comfit_chatbot = ComFitChatbot()
+    ui = ChatbotUI(comfit_chatbot)
+    demo = ui.create_ui()
+    demo.queue(max_size=10, default_concurrency_limit=3)
+    demo.launch(allowed_paths=["logo.png"])
+

config.json

@@ -0,0 +1 @@
+{ "models": ["hf.co/JatinkInnovision/ComFit4:Q4_K_M"]}

requirements.txt

@@ -0,0 +1,18 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
+numpy==1.26.4
+ollama==0.3.3
+onnx==1.17.0
+gradio==5.16.0
+gradiologin
+gradio[oauth]
+huggingface-hub
+hf-transfer
+llama-index-core==0.11.18
+llama-index-embeddings-huggingface-optimum==0.2.0
+llama-index-llms-ollama==0.3.4
+llama-index-vector-stores-duckdb==0.2.0
+llama-index-llms-openai==0.2.5
+duckdb==0.10.3
+nest_asyncio
+torch==2.5.0+cu121
+pydantic==2.10.6

start.sh

@@ -0,0 +1,3 @@
+#!/bin/sh
+ollama serve > /dev/null 2>&1 & 
+sleep 10 && python3 app.py