Update app.py

app.py

@@ -4,8 +4,9 @@ import json
 import logging
 import uuid
 import streamlit as st
+
 sys.path.append(os.path.abspath("src"))
-from src.chatbot import generate_response
+from src.chatbot import generate_response, COUNTRY_HELPLINES, get_helpline_for_country
 
 # Set up logging for safeguard triggers
 logging.basicConfig(filename="safeguard_logs.txt", level=logging.INFO)
@@ -13,66 +14,70 @@ logging.basicConfig(filename="safeguard_logs.txt", level=logging.INFO)
 # Persistent storage for metrics
 METRICS_FILE = "user_metrics.json"
 
-# Load existing metrics or initialize an empty structure
 def load_metrics():
+    """Load existing user metrics or initialize them."""
     if os.path.exists(METRICS_FILE):
         with open(METRICS_FILE, "r") as f:
             return json.load(f)
     return {"sessions": [], "feedback": {"thumbs_up": 0, "thumbs_down": 0}, "harmful_content": []}
 
-# Save metrics to file
 def save_metrics(metrics):
+    """Save updated user metrics to file."""
     with open(METRICS_FILE, "w") as f:
         json.dump(metrics, f, indent=4)
 
-# Initialize metrics
+# Initialize session
 if "metrics" not in st.session_state:
     st.session_state.metrics = load_metrics()
 
-# Initialize session ID
 if "session_id" not in st.session_state:
-    st.session_state.session_id = str(uuid.uuid4())[:8]  # Short unique ID for the session
+    st.session_state.session_id = str(uuid.uuid4())[:8]  # Short unique session ID
 
-# Streamlit app
+# Streamlit UI
 st.title("Fine-Tuned GPT Chatbot 🤖")
 
-# Chat history
+# --- Disclaimer ---
+st.markdown(
+    "**Disclaimer:** This chatbot is a prototype and is **not a substitute** for professional mental health advice. "
+    "If you are in crisis, please reach out to a professional or a mental health helpline."
+)
+
+# --- Country Selection ---
+st.subheader("Select Your Country for Helpline Information:")
+if "selected_country" not in st.session_state:
+    st.session_state.selected_country = "Other"  # Default
+
+st.session_state.selected_country = st.selectbox(
+    "Choose your country:", options=list(COUNTRY_HELPLINES.keys()), index=0
+)
+
+# --- Chat History ---
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
-# User input
+# --- User Input ---
 user_input = st.text_input("Ask me anything:")
 
 if st.button("Send"):
-    if user_input:
-        # Check for harmful content using the model's understanding
-        harmful_response = generate_response(
-            f"Detect if the following text contains harmful content (e.g., racism, violence, self-harm): {user_input}. "
-            "Respond with 'Harmful' if it does, otherwise respond with 'Safe'."
-        )
-
-        if "Harmful" in harmful_response:
-            # Log harmful content
-            st.session_state.metrics["harmful_content"].append({
-                "session_id": st.session_state.session_id,
-                "input": user_input,
-                "timestamp": st.session_state.get("timestamp", "N/A")
-            })
-            save_metrics(st.session_state.metrics)
-            st.warning("Your input contains harmful content and cannot be processed.")
-        else:
-            # Generate a response
-            response = generate_response(user_input)
-
-            # Store conversation
-            st.session_state.chat_history.append(f"You: {user_input}")
-            st.session_state.chat_history.append(f"Bot: {response}")
-
-# Display chat history
-for msg in st.session_state.chat_history:
-    st.write(msg)
-
-# Feedback system
+    if user_input.strip():
+        # Generate chatbot response using country selection
+        response = generate_response(user_input, st.session_state.selected_country)
+
+        # Store conversation
+        st.session_state.chat_history.append(f"You: {user_input}")
+        st.session_state.chat_history.append(f"Bot: {response}")
+
+        # Display chat history
+        for msg in st.session_state.chat_history:
+            st.write(msg)
+
+        # Display country-specific helpline
+        st.markdown(f"📞 **Helpline for {st.session_state.selected_country}:** {get_helpline_for_country(st.session_state.selected_country)}")
+
+    else:
+        st.warning("Please enter a message.")
+
+# --- Feedback System ---
 st.markdown("---")
 st.subheader("Feedback")
 col1, col2 = st.columns(2)
@@ -89,7 +94,7 @@ with col2:
         save_metrics(st.session_state.metrics)
         st.error("We're sorry to hear that. We'll try to improve! 😔")
 
-# User experience bar
+# --- User Experience Progress Bar ---
 total_feedback = (
     st.session_state.metrics["feedback"]["thumbs_up"] + st.session_state.metrics["feedback"]["thumbs_down"]
 )
@@ -106,6 +111,6 @@ else:
     st.subheader("User Experience")
     st.write("No feedback yet. Be the first to share your experience!")
 
-# Display session ID
+# --- Session ID Display ---
 st.markdown("---")
-st.write(f"Session ID: {st.session_state.session_id}")
+st.write(f"Session ID: {st.session_state.session_id}")