Update app.py

app.py

@@ -385,15 +385,26 @@ def load_user_data_source(user_id):
         st.error(f"Error loading user data source: {e}")
         return []
 
+def update_message_counter():
+    remaining_messages = st.session_state["message_limit"] - st.session_state["used_messages"]
+    message_counter_placeholder = st.sidebar.empty()
+
+    message_counter_placeholder.markdown(f"**Messages Left**: {remaining_messages} / {st.session_state['message_limit']}")
+
 
 
 def store_brand_tonality(user_id, message):
     try:
         tonality_id = str(uuid.uuid4())
-        db.child(f"users/{user_id}/brandTonality").child(tonality_id).set({"message": message})
-        if "brand_tonality" not in st.session_state:
-            st.session_state["brand_tonality"] = {}
-        st.session_state["brand_tonality"][tonality_id] = {"message": message}
+        # Save to Firebase
+        db.child("users").child(user_id).child("BrandTonality").child(tonality_id).set({"message": message})
+
+        # Update `st.session_state` for immediate sidebar display
+        if "BrandTonality" not in st.session_state:
+            st.session_state["BrandTonality"] = {}
+        st.session_state["BrandTonality"][tonality_id] = {"message": message}
+
+        # Confirmation
         display_save_confirmation("Brand Tonality")
 
     except Exception as e:
@@ -403,10 +414,15 @@ def store_brand_tonality(user_id, message):
 def store_trustbuilder(user_id, message):
     try:
         trustbuilder_id = str(uuid.uuid4())
-        db.child(f"users/{user_id}/trustBuilders").child(trustbuilder_id).set({"message": message})
-        if "trustbuilders" not in st.session_state:
-            st.session_state["trustbuilders"] = {}
-        st.session_state["trustbuilders"][trustbuilder_id] = {"message": message}
+        # Save to Firebase
+        db.child("users").child(user_id).child("TrustBuilder").child(trustbuilder_id).set({"message": message})
+
+        # Update `st.session_state` for immediate sidebar display
+        if "TrustBuilder" not in st.session_state:
+            st.session_state["TrustBuilder"] = {}
+        st.session_state["TrustBuilder"][trustbuilder_id] = {"message": message}
+
+        # Confirmation
         display_save_confirmation("TrustBuilder")
 
     except Exception as e:
@@ -414,49 +430,65 @@ def store_trustbuilder(user_id, message):
 
 
 def load_user_content(user_id):
+    """
+    Load all content for a user from Firebase, ensuring each user has a single root
+    containing TrustBuilder, BrandTonality, and other data fields like email, message limits, etc.
+    """
     try:
-        user_data = db.child(f"users/{user_id}").get().val()
+        user_data = db.child("users").child(user_id).get().val()
         if user_data:
+            # Update session state with all user data
             st.session_state.update(user_data)
-            st.session_state["trustbuilders"] = user_data.get("trustBuilders", {})
-            st.session_state["brand_tonality"] = user_data.get("brandTonality", {})
+            
+            # Load TrustBuilder and BrandTonality into session state for display
+            st.session_state["TrustBuilder"] = user_data.get("TrustBuilder", {})
+            st.session_state["BrandTonality"] = user_data.get("BrandTonality", {})
+            
     except Exception as e:
-        st.error(f"Error loading saved content: {e}")
         
         
 def save_content(user_id, content, content_type):
+    """
+    Save TrustBuilder or Brand Tonality to Firebase and update session state, with a clean confirmation.
+    """
     try:
         entry_id = str(uuid.uuid4())
-        db.child("users").child(user_id).child(content_type).child(entry_id).set({"message": content})
-        
-        # Update session state directly after saving to avoid duplicates
-        st.session_state[content_type.lower()][entry_id] = {"message": content}
-        display_save_confirmation(content_type)
+        firebase_key = "TrustBuilder" if content_type.lower() == "trustbuilder" else "BrandTonality"
+
+        # Save to Firebase
+        db.child("users").child(user_id).child(firebase_key).child(entry_id).set({"message": content})
+
+        # Update `st.session_state` for immediate display in sidebar and chatbot
+        if firebase_key not in st.session_state:
+            st.session_state[firebase_key] = {}
+        st.session_state[firebase_key][entry_id] = {"message": content}
+
+        # Display save confirmation without showing any ID or debug information
+        display_save_confirmation(content_type.capitalize())
+
     except Exception as e:
-        logger.error(f"Error saving {content_type}: {e}")
-        st.error(f"Could not save {content_type}. Please try again.")
+        st.error(f"Could not save {content_type}. Please try again. Error: {e}")
+
         
 def display_saved_trustbuilders():
-    """
-    Display saved TrustBuilders and Brand Tonality entries with download and delete options.
-    """
+    
+    # Display TrustBuilders with delete buttons
     st.sidebar.subheader("Saved TrustBuilders")
-    trustbuilders = st.session_state.get("trustbuilders", {})
+    trustbuilders = st.session_state.get("TrustBuilder", {})
     for trustbuilder_id, trustbuilder_data in trustbuilders.items():
-        st.sidebar.markdown(f"- {trustbuilder_data['message'][:50]}...")
-        st.sidebar.markdown(download_link(trustbuilder_data["message"], f"TrustBuilder_{trustbuilder_id}.md"), unsafe_allow_html=True)
-        if st.sidebar.button("❌", key=f"del_tb_{trustbuilder_id}"):
-            delete_entry("trustBuilders", trustbuilder_id, "trustbuilders")
-            st.rerun() 
+        message = trustbuilder_data.get("message", "No message available")
+        st.sidebar.markdown(f"- {message[:50]}...")
+        if st.sidebar.button("Delete", key=f"del_tb_{trustbuilder_id}"):
+            delete_entry("TrustBuilder", trustbuilder_id, "TrustBuilder")
 
+    # Display Brand Tonality entries with delete buttons
     st.sidebar.subheader("Saved Brand Tonality")
-    brand_tonality = st.session_state.get("brand_tonality", {})
+    brand_tonality = st.session_state.get("BrandTonality", {})
     for tonality_id, tonality_data in brand_tonality.items():
-        st.sidebar.markdown(f"- {tonality_data['message'][:50]}...")
-        st.sidebar.markdown(download_link(tonality_data["message"], f"BrandTonality_{tonality_id}.md"), unsafe_allow_html=True)
-        if st.sidebar.button("❌", key=f"del_bt_{tonality_id}"):
-            delete_entry("brandTonality", tonality_id, "brand_tonality")
-            st.rerun()
+        message = tonality_data.get("message", "No message available")
+        st.sidebar.markdown(f"- {message[:50]}...")
+        if st.sidebar.button("Delete", key=f"del_bt_{tonality_id}"):
+            delete_entry("BrandTonality", tonality_id, "BrandTonality")
 
 
 def download_link(content, filename):
@@ -527,13 +559,12 @@ def side():
         # Save Custom Content section
         st.header("Save Custom Content")
         custom_content = st.text_area("Enter content to save")
-        
         if st.button("Save as Brand Tonality"):
-            save_content(st.session_state["wix_user_id"], custom_content, "BrandTonality")
-            display_save_confirmation("Brand Tonality")
+            save_content(st.session_state["wix_user_id"], custom_content, "brandtonality")
         if st.button("Save as TrustBuilder"):
-            save_content(st.session_state["wix_user_id"], custom_content, "TrustBuilder")
-            display_save_confirmation("TrustBuilder")
+            save_content(st.session_state["wix_user_id"], custom_content, "trustbuilder")
+
+
 
 
 
@@ -1137,15 +1168,22 @@ if "brand_tonality" not in st.session_state:
 
 # Load saved entries upon user login
 def retrieve_user_data(user_id):
+    """
+    Load all content for a user from Firebase, ensuring each user has a single root
+    containing TrustBuilder, BrandTonality, and other data fields like email, message limits, etc.
+    """
     try:
-        trustbuilders = db.child(f"food_for_the_hungry/users/{user_id}/trustBuilders").get().val() or {}
-        brand_tonality = db.child(f"food_for_the_hungry/users/{user_id}/brandTonality").get().val() or {}
-        
-        # Set retrieved data into session state
-        st.session_state["trustbuilders"] = trustbuilders
-        st.session_state["brand_tonality"] = brand_tonality
+        user_data = db.child("users").child(user_id).get().val()
+        if user_data:
+            # Update session state with all user data
+            st.session_state.update(user_data)
+            
+            # Load TrustBuilder and BrandTonality into session state for display
+            st.session_state["TrustBuilder"] = user_data.get("TrustBuilder", {})
+            st.session_state["BrandTonality"] = user_data.get("BrandTonality", {})
+            
     except Exception as e:
-        st.error("Error loading saved entries. Please try again later.")
+        st.error(f"Error loading saved content: {e}")
 
 
 
@@ -1175,47 +1213,63 @@ def handle_memory_queries(prompt):
 
     # Retrieve saved TrustBuilders
     elif "show my saved trustbuilders" in prompt or "find my saved trustbuilders" in prompt:
-        if st.session_state.get("trustbuilders"):
-            return "\n".join([f"- {entry['message']}" for entry in st.session_state["trustbuilders"].values()])
+        trustbuilders = st.session_state.get("TrustBuilder", {})
+        if trustbuilders:
+            # Return only the content messages without any IDs or extra details
+            response = "\n".join([f"- {entry['message']}" for entry in trustbuilders.values()])
+            return f"Here are your saved TrustBuilders:\n{response}"
         else:
             return "You haven't saved any TrustBuilders yet."
 
     # Retrieve saved Brand Tonality
     elif "show my saved brand tonality" in prompt or "find my saved brand tonality" in prompt:
-        if st.session_state.get("brand_tonality"):
-            return "\n".join([f"- {entry['message']}" for entry in st.session_state["brand_tonality"].values()])
+        brand_tonality = st.session_state.get("BrandTonality", {})
+        if brand_tonality:
+            # Return only the content messages without any IDs or extra details
+            response = "\n".join([f"- {entry['message']}" for entry in brand_tonality.values()])
+            return f"Here are your saved Brand Tonality entries:\n{response}"
         else:
             return "You haven't saved any Brand Tonality entries yet."
-    
-    return None  # If no recognized command, proceed with general handling
+
+    return None # If no recognized command, proceed with general handling
 
 
 def delete_entry(category, entry_id, session_key):
+    """
+    Deletes an entry from Firebase and updates the session state to reflect the deletion.
+    """
     try:
+        # Get user ID from session state
         user_id = st.session_state["wix_user_id"]
+        
+        # Remove the entry from Firebase
         db.child("users").child(user_id).child(category).child(entry_id).remove()
-        del st.session_state[session_key][entry_id]  # Remove from session
-        st.experimental_rerun()  # Refresh sidebar
+
+        # Remove from session state and refresh sidebar display
+        if session_key in st.session_state and entry_id in st.session_state[session_key]:
+            del st.session_state[session_key][entry_id]
+
+        # Refresh the app to immediately reflect the change in the sidebar
+        st.rerun()
     except Exception as e:
         st.error(f"Error deleting entry: {e}")
 
+        
 # Function to download TrustBuilder as a .md file
 def download_trustbuilder_as_md(content, trustbuilder_id):
     b64_content = base64.b64encode(content.encode()).decode()
     download_link = f'<a href="data:text/markdown;base64,{b64_content}" download="TrustBuilder_{trustbuilder_id}.md">Download</a>'
     st.sidebar.markdown(download_link, unsafe_allow_html=True)
 
-
+def handle_save_trustbuilder_or_tonality(user_id, content, type_name):
+    save_content(user_id, content, type_name)
+    display_saved_trustbuilders()
 
 
                         
 # Function to update the message counter in a static location
-message_counter_placeholder = st.sidebar.empty()
 
 
-def update_message_counter():
-    remaining_messages = st.session_state["message_limit"] - st.session_state["used_messages"]
-    message_counter_placeholder.markdown(f"**Messages Left**: {remaining_messages} / {st.session_state['message_limit']}")
 
 def initialize_wix_only_user(wix_user_id, email):
     try:
@@ -1288,11 +1342,11 @@ if st.session_state.get("wix_user_id") and st.session_state.get("email"):
                 if save_as_trustbuilder or save_as_tonality:
                     user_id = st.session_state["wix_user_id"]
                     if save_as_trustbuilder:
-                        store_trustbuilder(user_id, prompt)
-                        display_save_confirmation("TrustBuilder")
+                        handle_save_trustbuilder_or_tonality(user_id, prompt, "trustbuilder")
                     elif save_as_tonality:
-                        store_brand_tonality(user_id, prompt)
-                        display_save_confirmation("Brand Tonality")
+                        handle_save_trustbuilder_or_tonality(user_id, prompt, "brandtonality")
+
+                   
                 else:
                     # Generate a response with AI for general queries
                     with st.chat_message("user"):