Update src/session_state.py

src/session_state.py

@@ -35,10 +35,11 @@ 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:
-                db_history = st.session_state.db_client.get_conversation_history(st.session_state.consultation_id)
-                st.session_state.history = db_history
+                # Load initial history from DB
+                st.session_state.history = load_history_from_db(st.session_state.consultation_id)
             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
@@ -70,16 +71,20 @@ def initialize_session_state():
 
 def add_message_to_history(message):
     """
-    Add a message to the conversation history and persist it to database.
+    Add a message to the local 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:
+    if hasattr(st.session_state, 'db_client') and st.session_state.db_client and hasattr(st.session_state, 'consultation_id'):
         try:
             st.session_state.db_client.save_message(
                 st.session_state.consultation_id, 
@@ -88,55 +93,80 @@ def add_message_to_history(message):
         except Exception as e:
             st.error(f"Failed to save message to database: {str(e)}")
 
-def get_full_history():
+def load_history_from_db(consultation_id):
     """
-    Get the complete conversation history from the database.
+    Load the complete conversation history from the database for a given consultation_id.
     
+    Args:
+        consultation_id (str): The ID of the consultation to load.
+        
     Returns:
-        list: Full conversation history
+        list: Conversation history from the database, or an empty list on failure.
     """
-    if not hasattr(st.session_state, 'consultation_id'):
-        print("No consultation_id in session state")
+    if not consultation_id:
+        print("No consultation_id provided for loading history")
         return []
         
     if not hasattr(st.session_state, 'db_client') or not st.session_state.db_client:
-        print("No database client available")
-        return st.session_state.history if hasattr(st.session_state, 'history') else []
+        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(
-            st.session_state.consultation_id
+            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}")
-        
+        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)}")
-        
-        # Fallback to session state history if database retrieval fails
-        return st.session_state.history if hasattr(st.session_state, 'history') else []
+        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', [])
+    
+    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}")
+        return db_history
+    except Exception as e:
+        print(f"Failed to sync history from database: {str(e)}")
+        return st.session_state.get('history', [])
+
 
 def end_conversation():
     """
     End the current conversation and clean up resources.
     """
-    if hasattr(st.session_state, 'db_client') and st.session_state.db_client:
+    if hasattr(st.session_state, 'db_client') and st.session_state.db_client and hasattr(st.session_state, 'consultation_id'):
         try:
-            st.session_state.db_client.delete_conversation(
-                st.session_state.consultation_id
-            )
+            # 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
         except Exception as e:
-            st.error(f"Failed to delete conversation from database: {str(e)}")
+            st.error(f"Failed to update conversation status in database: {str(e)}")
     
-    # Reset session state
+    # Reset session state for a new conversation
     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: