Update app.py

app.py

@@ -17,7 +17,6 @@ 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
 
@@ -176,7 +175,6 @@ class VectorStoreManager:
         """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
     
@@ -201,15 +199,12 @@ vector_store_manager = VectorStoreManager()
 
 
 
-
-
 class DEKCIBChatbot:
     def __init__(self):
         self.initialize()
 
 
     def initialize(self):
-        # self.user_id = self.get_user_data(USER_NAME)        
         self.session_manager = SessionManager()
         self.embed_model = OptimumEmbedding(folder_name=EMBED_MODEL_PATH)
         Settings.embed_model = self.embed_model
@@ -218,11 +213,6 @@ class DEKCIBChatbot:
         self.config = self._load_config()
         self.llm_options = self._initialize_models()
 
-        # default_vector_store = self.vector_stores[DEFAULT_VECTOR_STORE]["data"]
-        # self.session_manager.create_session(USER_NAME)
-
-
-
 
 
     def get_user_data(self, user_id):
@@ -347,9 +337,6 @@ class DeKCIBChatEngine:
         self.chat_engine = None
         self.rebuild_chat_engine_flag = True
 
-
-        # self.llm_name = DEFAULT_LLM
-        # self.vector_store_name = DEFAULT_VECTOR_STORE
         
         # Conversation metadata management
         self.convs_metadata = {}
@@ -381,10 +368,7 @@ class DeKCIBChatEngine:
             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()
-        # else:
-        #     self.build_convs_metadata()
-        #     with open(self.convs_metadata_file_path, "w") as f:
-        #         json.dump(self.convs_metadata, f)
+
 
 
 
@@ -452,30 +436,17 @@ class DeKCIBChatEngine:
         )
         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"""
@@ -517,7 +488,7 @@ class DeKCIBChatEngine:
         
 
         print(f"Initializing memory for conversation {conversation_id}")
-
+        print(self.chat_store)
 
         self.memory = ChatMemoryBuffer.from_defaults(
             token_limit=3000,
@@ -546,8 +517,6 @@ class DeKCIBChatEngine:
             system_prompt=SYSTEM_PROMPT
         )
 
-
-
         self.set_rebuild_chat_engine_flag(False)
         return self.chat_engine
     
@@ -570,55 +539,12 @@ class DeKCIBChatEngine:
         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)
@@ -635,27 +561,31 @@ class DeKCIBChatEngine:
             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
-
-        elapsed_time = time.time() - start_time
-
+        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)
 
-        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)
-
+        print("update_convs_metadata")
         self.save_metadata()
-
+        print("save_metadata")
         self.save_chat_history()
-
+        print("save_chat_history")
         
         return response
     
@@ -744,7 +674,34 @@ class DeKCIBChatEngine:
         """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"""
@@ -765,56 +722,46 @@ class DeKCIBChatEngine:
                 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.update_convs_metadata(conversation_id)
-            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.update_convs_metadata(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)
-
-            
+    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)
             
-
-            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()
+        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
 
 
 
@@ -850,6 +797,23 @@ class ChatbotUI:
     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
+
+
 
 
 
@@ -882,7 +846,7 @@ class ChatbotUI:
                 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()],
+                        choices=[(v["display_name"], k) for k, v in self.vector_stores.items()],
                         value=next(iter(self.vector_stores.keys()), None)
 
                     )
@@ -909,7 +873,7 @@ class ChatbotUI:
                         type="index",
                         layout="table"
                     )
-                    
+
                 
                 # Main chat area
                 with gr.Column(scale=3):
@@ -945,74 +909,62 @@ class ChatbotUI:
 
 
 
+
             def chat_with_dekcib(history, user_id, conv_idx):
                 
-                start_time = time.time()
 
                 msg = history[-1]["content"]
 
                 user_engine = self.dekcib_chatbot.session_manager.sessions[user_id]
                 # user_engine.che
                 
-                # conv_id = None
+
+                print("conv_idx")
+                print(conv_idx)
+                
+
+                conv_id = None
                 if conv_idx is not None:
                     conv_id = user_engine.sorted_conversation_list[conv_idx]
-                else:
+
+                if len(history) == 1:
                     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)
+                    history.append({"role": "assistant", "content": "", "metadata":{"Thinking...": ""}})
                     for character in think_content:
-                        history[-1]["content"] += character
+                        history[-1]["metadata"]["Thinking..."] += 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():
+                print("clear_msg")
                 return ""
             
 
@@ -1049,32 +1001,47 @@ class ChatbotUI:
 
 
             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)
 
-                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)
-
-            
+                return gr.update(value=conversation_history), chat_history
 
             
-
-
-
             conversation_history.click(
                 click_to_select_conversation,
                 [conversation_history, user_id],
-                [conversation_history, chatbot, llm_dropdown, vector_dropdown]
+                [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:
@@ -1084,11 +1051,21 @@ class ChatbotUI:
                 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
@@ -1097,10 +1074,8 @@ class ChatbotUI:
                 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)
+                chat_history = user_engine.get_chat_history(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"]
 
                 yield llm_name, vector_store_name, update_conversation_history, chat_history
 
@@ -1124,7 +1099,6 @@ class ChatbotUI:
                 user_engine.set_llm(llm_name)
 
 
-
             llm_dropdown.change(
                 update_llm,
                 [user_id, llm_dropdown],
@@ -1146,119 +1120,8 @@ class ChatbotUI:
                 [user_id, vector_dropdown],
                 None
             )
-
-
-            def edit_chat(user_id, history, edit_data: EditData):
-                user_engine = self.dekcib_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_dekcib,
-                    [chatbot, user_id, conversation_history],
-                    [chatbot]
-                ).success(
-                    update_conversation_history,
-                    [user_id],
-                    [conversation_history]
-                )
-
-
-
-            def retry_chat(user_id, history):
-                user_engine = self.dekcib_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_dekcib,
-                    [chatbot, user_id, conversation_history],
-                    [chatbot]
-                ).then(
-                    update_conversation_history,
-                    [user_id],
-                    [conversation_history]
-                )
-
-
-            def undo_chat(user_id):
-                user_engine = self.dekcib_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.dekcib_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():
@@ -1267,6 +1130,7 @@ class ChatbotUI:
                 return [], gr.update(value=None)
             
             def print_dataset(value):
+                print("value")
                 print(value)
             
             # Create new conversation
@@ -1286,11 +1150,17 @@ class ChatbotUI:
 
 # 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"])
+