Update app.py

app.py

@@ -1569,30 +1569,27 @@ def retrieve_user_data(user_id):
         st.error(f"Error loading saved content: {e}")
 
 
-
-
-
-# Function to handle queries about saved TrustBuilders or Brand Tonality
 def handle_memory_queries(prompt):
     prompt = prompt.lower().strip()
     
-    # Save as TrustBuilder, removing the phrase from the content
-    if "save this as trustbuilder" in prompt:
-        content_to_save = prompt.replace("save this as trustbuilder", "").strip()
-        if content_to_save:  # Only save if there's actual content
-            save_content(st.session_state["wix_user_id"], content_to_save)
-            return "TrustBuilder saved successfully!"
+    # Save as TrustBuilder
+    if re.search(r"\b(keep|add|save|store)\s*(this)?\s*(as)?\s*(trust\s*builder|trustbuilder)\b", prompt, re.IGNORECASE):
+        content_to_save = re.sub(r"\b(keep|add|save|store)\s*(this)?\s*(as)?\s*(trust\s*builder|trustbuilder)\b", "", prompt, flags=re.IGNORECASE).strip()
+        specified_bucket = None
+
+        # Check for explicit bucket mention in the same prompt
+        match = re.search(r"under\s+(\w+\s*trust)", prompt, re.IGNORECASE)
+        if match:
+            bucket_name = match.group(1).replace("trust", "").strip()
+            if bucket_name in ["Stability", "Development", "Relationship", "Benefit", "Vision", "Competence"]:
+                specified_bucket = bucket_name
+
+        if content_to_save:
+            handle_save_trustbuilder(content_to_save, specified_bucket)
         else:
-            return "Please provide content to save as TrustBuilder."
-
-    # Save as Brand Tonality, removing the phrase from the content
-    elif "save this as brand tonality" in prompt:
-        content_to_save = prompt.replace("save this as brand tonality", "").strip()
-        if content_to_save:  # Only save if there's actual content
-            save_content(st.session_state["wix_user_id"], content_to_save)
-            return "Brand Tonality saved successfully!"
-        else:
-            return "Please provide content to save as Brand Tonality."
+            with st.chat_message("assistant"):
+                st.markdown("Please provide content to save as TrustBuilder.")
+        return None
 
     # Retrieve saved TrustBuilders
     elif "show my saved trustbuilders" in prompt or "find my saved trustbuilders" in prompt:
@@ -1614,9 +1611,8 @@ def handle_memory_queries(prompt):
         else:
             return "You haven't saved any Brand Tonality entries yet."
 
-    return None # If no recognized command, proceed with general handling
-
-
+    return None  # 
+    
 def delete_entry(category, entry_id):
     try:
         user_id = st.session_state["wix_user_id"]
@@ -1648,17 +1644,20 @@ def handle_save_trustbuilder(content, specified_bucket=None):
                 break
 
     if not bucket:
-        st.warning("No Trust Bucket specified. Please indicate the Trust Bucket.")
+        # Store content and prompt user for bucket
         st.session_state["missing_trustbucket_content"] = content
+        with st.chat_message("assistant"):
+            st.markdown(
+                "No Trust Bucket specified. Please indicate the Trust Bucket (e.g., Stability, Development, Relationship, Benefit, Vision, Competence)."
+            )
         return
 
     # Save TrustBuilder with detected/provided bucket
     brand = st.session_state.get("brand_input_save", "Unknown")
     content_to_save = f"{bucket}: Brand: {brand.strip()} | {content.strip()}"
     save_content(st.session_state["wix_user_id"], content_to_save)
-    st.success(f"TrustBuilder saved under '{bucket}' Trust.")
-
-
+    with st.chat_message("assistant"):
+        st.markdown(f"TrustBuilder saved under '{bucket}' Trust.")
 
                         
 # Function to update the message counter in a static location
@@ -1672,6 +1671,9 @@ if "message_limit" not in st.session_state:
 if "used_messages" not in st.session_state:
     st.session_state["used_messages"] = 0
 
+if "missing_trustbucket_content" not in st.session_state:
+    st.session_state["missing_trustbucket_content"] = None
+
 def initialize_user_session():
     """
     Initialize user session and ensure user data exists in Firebase.
@@ -1730,42 +1732,23 @@ if st.session_state["used_messages"] < st.session_state["message_limit"]:
     if prompt:
         st.session_state.chat_started = True
         st.session_state.chat_history.append({"role": "user", "content": prompt})
-        if "save this as trustbuilder" in prompt.lower():
-            content_to_save = prompt.replace("save this as trustbuilder", "").strip()
-            specified_bucket = None
-
-            # Check for explicit bucket mention in the same prompt
-            match = re.search(r"under\s+(\w+\s*trust)", prompt, re.IGNORECASE)
-            if match:
-                bucket_name = match.group(1).replace("trust", "").strip()
-                if bucket_name in ["Stability", "Development", "Relationship", "Benefit", "Vision", "Competence"]:
-                    specified_bucket = bucket_name
-
-            if content_to_save:
-                handle_save_trustbuilder(content_to_save, specified_bucket)
-            else:
-                st.warning("Please provide content to save as TrustBuilder.")
-
-
+  
 
         # Handle save commands based on the user prompt
         memory_response = handle_memory_queries(prompt)
         if memory_response:
             with st.chat_message("assistant"):
                 st.markdown(memory_response)
-        else:
-            save_as_trustbuilder = re.search(r"\b(save|add|store)\s*(this)?\s*(as)?\s*(trust\s*builder|trustbuilder)\b", prompt, re.IGNORECASE)
-            save_as_tonality = re.search(r"\b(save|add|store)\s*(this)?\s*(as)?\s*(brand\s*tonality|tonality)\b", prompt, re.IGNORECASE)
-
-            if save_as_trustbuilder or save_as_tonality:
-                user_id = st.session_state["wix_user_id"]
-                if save_as_trustbuilder:
-                    handle_save_trustbuilder_or_tonality(user_id, prompt, "trustbuilder")
-                elif save_as_tonality:
-                    handle_save_trustbuilder_or_tonality(user_id, prompt, "brandtonality")
-
-               
+        elif st.session_state.get("missing_trustbucket_content"):
+            bucket = prompt.strip()
+            valid_buckets = ["Stability", "Development", "Relationship", "Benefit", "Vision", "Competence"]
+            if bucket in valid_buckets:
+                content_to_save = st.session_state.pop("missing_trustbucket_content")
+                handle_save_trustbuilder(content_to_save, bucket)
             else:
+                with st.chat_message("assistant"):
+                    st.markdown("Invalid Trust Bucket. Please provide a valid Trust Bucket (e.g., Stability, Development, Relationship, Benefit, Vision, Competence).")
+        else:
                 # Generate a response with AI for general queries
                 with st.chat_message("user"):
                     st.markdown(prompt)