Update src/session_state.py

src/session_state.py

@@ -35,11 +35,10 @@ def initialize_session_state():
         # Try to load from database if consultation_id exists
         if hasattr(st.session_state, 'consultation_id') and st.session_state.db_client:
             try:
-                # Load initial history from DB
-                st.session_state.history = load_history_from_db(st.session_state.consultation_id)
+                db_history = st.session_state.db_client.get_conversation_history(st.session_state.consultation_id)
+                st.session_state.history = db_history
             except Exception as e:
                 # Silently continue with empty history if database fetch fails
-                print(f"Initial history load failed: {str(e)}")
                 pass
     
     # RAG feature toggle
@@ -71,20 +70,16 @@ def initialize_session_state():
 
 def add_message_to_history(message):
     """
-    Add a message to the local conversation history and persist it to database.
+    Add a message to the conversation history and persist it to database.
     
     Args:
         message (dict): The message to add to history
     """
-    # Ensure history list exists
-    if 'history' not in st.session_state:
-        st.session_state.history = []
-        
     # Add to local history
     st.session_state.history.append(message)
     
     # Persist to database if available
-    if hasattr(st.session_state, 'db_client') and st.session_state.db_client and hasattr(st.session_state, 'consultation_id'):
+    if hasattr(st.session_state, 'db_client') and st.session_state.db_client:
         try:
             st.session_state.db_client.save_message(
                 st.session_state.consultation_id, 
@@ -93,80 +88,55 @@ def add_message_to_history(message):
         except Exception as e:
             st.error(f"Failed to save message to database: {str(e)}")
 
-def load_history_from_db(consultation_id):
+def get_full_history():
     """
-    Load the complete conversation history from the database for a given consultation_id.
+    Get the complete conversation history from the database.
     
-    Args:
-        consultation_id (str): The ID of the consultation to load.
-        
     Returns:
-        list: Conversation history from the database, or an empty list on failure.
+        list: Full conversation history
     """
-    if not consultation_id:
-        print("No consultation_id provided for loading history")
+    if not hasattr(st.session_state, 'consultation_id'):
+        print("No consultation_id in session state")
         return []
         
     if not hasattr(st.session_state, 'db_client') or not st.session_state.db_client:
-        print("No database client available for loading history")
-        return [] # Cannot load from DB, return empty
-    
-    try:
-        db_history = st.session_state.db_client.get_conversation_history(
-            consultation_id
-        )
-        print(f"Retrieved {len(db_history)} messages from database for consultation {consultation_id}")
-        return db_history
-    except Exception as e:
-        print(f"Failed to retrieve history from database: {str(e)}")
-        return [] # Return empty list on failure
-
-# Renamed get_full_history to reflect its purpose (it's now primarily for external use if needed)
-# The main way to get history for display is directly from st.session_state.history
-def get_synced_history_from_db():
-    """
-    Gets the history from the database and updates the local cache.
-    This is for ensuring sync if other processes might have updated the DB.
-    For most UI purposes, st.session_state.history should be used directly.
-    """
-    if not hasattr(st.session_state, 'consultation_id'):
-        print("No consultation_id in session state for sync")
-        return st.session_state.get('history', [])
-
-    if not hasattr(st.session_state, 'db_client') or not st.session_state.db_client:
-        print("No database client available for sync")
-        return st.session_state.get('history', [])
+        print("No database client available")
+        return st.session_state.history if hasattr(st.session_state, 'history') else []
     
     try:
         db_history = st.session_state.db_client.get_conversation_history(
             st.session_state.consultation_id
         )
-        st.session_state.history = db_history # Sync local cache
-        print(f"Synced history: Retrieved {len(db_history)} messages from database for consultation {st.session_state.consultation_id}")
+        
+        # Update the local cache with the database results
+        st.session_state.history = db_history
+        
+        # For debugging
+        print(f"Retrieved {len(db_history)} messages from database for consultation {st.session_state.consultation_id}")
+        
         return db_history
     except Exception as e:
-        print(f"Failed to sync history from database: {str(e)}")
-        return st.session_state.get('history', [])
-
+        print(f"Failed to retrieve history from database: {str(e)}")
+        
+        # Fallback to session state history if database retrieval fails
+        return st.session_state.history if hasattr(st.session_state, 'history') else []
 
 def end_conversation():
     """
     End the current conversation and clean up resources.
     """
-    if hasattr(st.session_state, 'db_client') and st.session_state.db_client and hasattr(st.session_state, 'consultation_id'):
+    if hasattr(st.session_state, 'db_client') and st.session_state.db_client:
         try:
-            # Potentially mark conversation as inactive instead of deleting
-            # st.session_state.db_client.delete_conversation(st.session_state.consultation_id)
-            pass # Placeholder if we decide not to delete immediately
+            st.session_state.db_client.delete_conversation(
+                st.session_state.consultation_id
+            )
         except Exception as e:
-            st.error(f"Failed to update conversation status in database: {str(e)}")
+            st.error(f"Failed to delete conversation from database: {str(e)}")
     
-    # Reset session state for a new conversation
+    # Reset session state
     st.session_state.history = []
-    old_consultation_id = st.session_state.get('consultation_id', 'N/A')
     st.session_state.consultation_id = str(uuid.uuid4())[:8]
-    print(f"Ended consultation {old_consultation_id}. Starting new consultation: {st.session_state.consultation_id}")
-
+    
     # Create new conversation in database
     if hasattr(st.session_state, 'db_client') and st.session_state.db_client:
         try: