Update app.py

app.py

@@ -436,16 +436,27 @@ 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):
@@ -537,6 +548,46 @@ class DeKCIBChatEngine:
         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):
@@ -544,7 +595,7 @@ class DeKCIBChatEngine:
 
     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)
@@ -561,31 +612,27 @@ 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
+        # 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)
+        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)
-        print("update_convs_metadata")
+
         self.save_metadata()
-        print("save_metadata")
+
         self.save_chat_history()
-        print("save_chat_history")
+
         
         return response
     
@@ -674,34 +721,7 @@ 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"""
@@ -722,47 +742,54 @@ class DeKCIBChatEngine:
                 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)
+    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)
+
+            
             
-        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
 
+            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()
 
 
 
@@ -846,7 +873,7 @@ class ChatbotUI:
                 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()],
+                        choices=[(v["display_name"]) for k, v in self.vector_stores.items()],
                         value=next(iter(self.vector_stores.keys()), None)
 
                     )
@@ -909,62 +936,74 @@ 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
                 
-
-                print("conv_idx")
-                print(conv_idx)
-                
-
-                conv_id = None
+                # conv_id = None
                 if conv_idx is not None:
                     conv_id = user_engine.sorted_conversation_list[conv_idx]
-
-                if len(history) == 1:
+                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"]
 
 
-                    history.append({"role": "assistant", "content": "", "metadata":{"Thinking...": ""}})
+                    # 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]["metadata"]["Thinking..."] += character
+                        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():
-                print("clear_msg")
                 return ""
             
 
@@ -1001,18 +1040,29 @@ 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)
 
-                return gr.update(value=conversation_history), chat_history
+                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]
+                [conversation_history, chatbot, llm_dropdown, vector_dropdown]
             )
 
 
@@ -1051,16 +1101,14 @@ 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."],
@@ -1074,8 +1122,10 @@ 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(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"]
 
                 yield llm_name, vector_store_name, update_conversation_history, chat_history
 
@@ -1099,6 +1149,7 @@ class ChatbotUI:
                 user_engine.set_llm(llm_name)
 
 
+
             llm_dropdown.change(
                 update_llm,
                 [user_id, llm_dropdown],
@@ -1120,8 +1171,168 @@ 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]
+                )
+
+
+            # def clear_chatbot(user_id):
+            #     print("user_id")
+            #     return []
             
+            # user_id_dropdown.change(
+            #     clear_chatbot,
+            #     [user_id_dropdown],
+            #     [chatbot]
+            # )
+
+            # # Load selected conversation from Dataset click
+            # # def load_clicked_conversation(evt: gr.SelectData, user_id, session_state):
+            # #     if not user_id or not evt.index:
+            # #         return [], session_state
+                    
+            # #     # Get all conversations for this user
+            # #     conversations = get_user_conversations(user_id)
+                
+            # #     # Check if we have enough conversations
+            # #     if not conversations or evt.index >= len(conversations):
+            # #         return [], session_state
+                    
+            # #     # Get the clicked conversation ID
+            # #     selected_conv_id = conversations[evt.index][0]  # Get the conversation ID from the tuple
+                
+            # #     # Load the conversation
+            # #     chat_store, chat_history = load_conversation(user_id, selected_conv_id)
+                
+            # #     if chat_store:
+            # #         # Update session state with loaded conversation
+            # #         if not session_state:
+            # #             session_state = {}
+                    
+            # #         session_state["user_id"] = user_id
+            # #         session_state["session_id"] = user_id
+            # #         session_state["conversation_id"] = selected_conv_id
+            # #         session_state["chat_store"] = chat_store
+            # #         session_state["total_score"] = 0.0
+            # #         session_state["answer_count"] = 0
+                    
+            # #         return chat_history, session_state
+            # #     else:
+            # #         return [], session_state
+            
+            # # Setup conversation history click handler
+            # conversation_history.select(
+            #     load_clicked_conversation,
+            #     [user_id_dropdown, session_state],
+            #     [chatbot, session_state]
+            # )
             
             # Create new conversation button should only clear the chat area, but not create a new conversation yet
             def prepare_new_chat():
@@ -1130,7 +1341,6 @@ class ChatbotUI:
                 return [], gr.update(value=None)
             
             def print_dataset(value):
-                print("value")
                 print(value)
             
             # Create new conversation