Update app.py

app.py

@@ -1621,29 +1621,46 @@ def retrieve_user_data(user_id):
         st.error(f"Error loading saved content: {e}")
 
 def handle_memory_queries(prompt):
+    """
+    Main function to handle user commands and allocate Trust Buckets.
+    """
     prompt = prompt.lower().strip()
 
     # Save and allocate TrustBuilder
     if "save and allocate" in prompt:
         content_to_save = prompt.replace("save and allocate:", "").strip()
         if content_to_save:
-            handle_save_trustbuilder(content_to_save)
+            assistant_response = handle_save_trustbuilder(content_to_save)
         else:
-            with st.chat_message("assistant"):
-                st.markdown("Please provide content to save and allocate.")
+            assistant_response = "Please provide content to save and allocate."
+
+        # Save response to chat history and display it
+        st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+        with st.chat_message("assistant"):
+            st.markdown(assistant_response)
         return None
 
     # Show saved TrustBuilders
-    elif "show my saved trustbuilders" in prompt or "find my saved trustbuilders" in prompt:
+    elif "find my saved trustbuilders" in prompt or "show my saved trustbuilders" in prompt:
         trustbuilders = fetch_trustbuilders(st.session_state.get("wix_user_id", "default_user"))
         if trustbuilders:
             saved_content = "\n".join([f"- {entry['message']}" for entry in trustbuilders.values()])
-            return f"Here are your saved TrustBuilders:\n{saved_content}"
+            assistant_response = f"Here are your saved TrustBuilders:\n{saved_content}"
         else:
-            return "You haven't saved any TrustBuilders yet."
+            assistant_response = "You haven't saved any TrustBuilders yet."
 
+        # Save response to chat history and display it
+        st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+        with st.chat_message("assistant"):
+            st.markdown(assistant_response)
+        return None
+
+    # Default case for generic AI response
+    assistant_response = "I'm sorry, I couldn't understand your request."
+    st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+    with st.chat_message("assistant"):
+        st.markdown(assistant_response)
     return None
-    
 def delete_entry(category, entry_id):
     try:
         user_id = st.session_state["wix_user_id"]
@@ -1660,17 +1677,36 @@ def download_trustbuilder_as_md(content, trustbuilder_id):
     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(content, specified_bucket=None):
     """
     Handles saving TrustBuilders by detecting or automatically allocating the Trust Bucket.
     """
     trust_buckets = {
-        "Stability": ["secure", "reliable", "consistent", "stable"],
-        "Development": ["growth", "future", "plan", "innovate", "strategy"],
-        "Relationship": ["trust", "partnership", "collaborate", "customer"],
-        "Benefit": ["value", "profit", "advantage", "benefit"],
-        "Vision": ["future", "goal", "mission", "aspire", "dream"],
-        "Competence": ["skill", "expertise", "knowledge", "competent", "ability"]
+        "Stability": [
+            "secure", "reliable", "consistent", "stable", "safety", "protection", "trustworthy",
+            "dependable", "solid", "durable", "continuity", "assurance"
+        ],
+        "Development": [
+            "growth", "future", "plan", "innovate", "strategy", "building", "improvement", "progress",
+            "expansion", "visionary", "change", "adaptation", "forward-thinking"
+        ],
+        "Relationship": [
+            "trust", "partnership", "collaborate", "customer", "engage", "connect", "loyalty",
+            "bond", "interaction", "relationship", "teamwork", "communication"
+        ],
+        "Benefit": [
+            "value", "profit", "advantage", "benefit", "gain", "cost-effective", "economical",
+            "saving", "reward", "output", "utility", "return", "efficiency"
+        ],
+        "Vision": [
+            "goal", "mission", "aspire", "dream", "visionary", "great", "future", "ideal",
+            "ambition", "long-term", "objective", "focus", "drive"
+        ],
+        "Competence": [
+            "skill", "expertise", "knowledge", "competent", "ability", "talent", "proficiency",
+            "capability", "qualification", "training", "effectiveness", "specialization"
+        ]
     }
 
     bucket = specified_bucket
@@ -1682,27 +1718,24 @@ def handle_save_trustbuilder(content, specified_bucket=None):
                 bucket = tb
                 break
 
-    # If still no bucket found, prompt the user
+    # If no bucket can be allocated, prompt the user
     if not bucket:
         st.session_state["missing_trustbucket_content"] = content
-        with st.chat_message("assistant"):
-            st.markdown(
-                "No Trust Bucket could be allocated automatically. Please indicate the Trust Bucket (e.g., Stability, Development, Relationship, Benefit, Vision, Competence)."
-            )
-        return
+        return (
+            "No Trust Bucket could be allocated automatically. "
+            "Please indicate the Trust Bucket (e.g., Stability, Development, Relationship, Benefit, Vision, Competence)."
+        )
 
     # 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)
 
-    # Confirm saving to the user
-    with st.chat_message("assistant"):
-        st.markdown(f"TrustBuilder allocated to **{bucket}** and saved successfully!")
-
+    # Return success response
+    return f"TrustBuilder allocated to **{bucket}** and saved successfully!"
 
 
-                        
+    
 # Function to update the message counter in a static location
 
 if "email" not in st.session_state: