# ============================================================
# FILE: app.py (Consolidated Root)
# PURPOSE: Single-file deployment for Safeguard AI.
# Contains UI, Navigation, Groq Bot Logic, and In-Memory Analytics.
# *UPDATED: Implements Lazy-Loading for SHAP XAI*
# *UPDATED: Per-browser session_id + LangChain LCEL LLM chain*
# ============================================================
import os
import sys
import uuid
import pandas as pd
import streamlit as st
import streamlit.components.v1 as components
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
# ─────────────────────────────────────────────────────────────
# PATH SETUP & BACKEND IMPORTS
# ─────────────────────────────────────────────────────────────
# Add project root to path so we can import the heavy ML engines from src/
project_root = os.path.dirname(os.path.abspath(__file__))
if project_root not in sys.path:
sys.path.append(project_root)
# Import the core AI engines (ensure your src/ folder is uploaded to HF!)
try:
from src.core_model.predict import MindGuardPredictor
from src.rag_engine.retriever import MindGuardRetriever
from src.audio.speech_to_text import MindGuardAudioProcessor
from src.explainability.shap_explainer import MindGuardSHAPExplainer
except ImportError as e:
st.error(f"Failed to import backend modules. Ensure the 'src' folder is uploaded. Details: {e}")
st.stop()
# ─────────────────────────────────────────────────────────────
# PAGE CONFIGURATION
# ─────────────────────────────────────────────────────────────
st.set_page_config(page_title="Mindguard AI", page_icon="🛡️", layout="wide")
# ─────────────────────────────────────────────────────────────
# CORE CHATBOT ENGINE
# ─────────────────────────────────────────────────────────────
class SafeguardChatbot:
"""Orchestrates prediction, semantic search, and Groq LLM response."""
def __init__(self):
# 1. Load API Key
load_dotenv(os.path.join(project_root, ".env"))
api_key = os.environ.get("GROQ_API_KEY")
if not api_key:
st.error("❌ GROQ_API_KEY not found in Hugging Face Secrets!")
st.stop()
# 2. Initialize the Groq LLM client through LangChain
self.llm = ChatGroq(
model="openai/gpt-oss-120b",
temperature=0.3,
api_key=api_key,
)
# 3. Wake up tools
self.predictor = MindGuardPredictor()
self.retriever = MindGuardRetriever()
self.audio_processor = MindGuardAudioProcessor()
self.system_prompt = """
You are Safeguard AI, a highly empathetic, clinical-grade mental health AI.
Your goal is to de-escalate emotional distress and provide actionable coping strategies.
STRICT RULES:
1. NEVER hallucinate medical advice. ONLY use the 'Clinical Strategy' provided in the prompt.
2. Keep your response conversational, warm, and easy to read (use short paragraphs).
3. Do not sound like a robot reading a textbook. Weave the clinical strategy naturally into your empathy.
4. If the Risk Level is 'High', prioritize grounding the user immediately.
"""
# 4. Assemble the LangChain LCEL generation chain
# System prompt + per-turn augmented prompt -> Groq LLM -> plain string
# prompt -> llm -> output_parser
self.prompt_template = ChatPromptTemplate.from_messages([
("system", self.system_prompt),
("human", "{augmented_prompt}"),
])
self.generation_chain = self.prompt_template | self.llm | StrOutputParser()
def generate_response(self, user_input, chat_history_list, session_id):
"""Pipeline: Predict -> Retrieve -> Remember -> Generate.
The caller must provide a session_id so every request is tied to one
browser session. Returns a dict with the reply plus all its metadata,
so the UI never has to guess or re-query anything.
"""
if not session_id:
raise ValueError("session_id is required for every chatbot request")
# 1. Core Model Prediction
prediction = self.predictor.predict(user_input)
emotion = prediction['emotion']
risk = prediction['risk_level']
# 2. RAG Retrieval
strategy = self.retriever.get_coping_strategy(
user_query=user_input,
emotion_filter=emotion
)
# 3. Extract recent history directly from Streamlit session state
history_text = "No previous conversation."
if chat_history_list:
recent = chat_history_list[-4:] # Grab last two exchanges
history_text = "\n".join([f"{msg['role'].capitalize()}: {msg['content']}" for msg in recent])
# 4. Prompt Engineering
augmented_prompt = f"""
--- RECENT CONVERSATION HISTORY ---
{history_text}
--- CURRENT SITUATION ---
User's New Message: "{user_input}"
AI Core Diagnosis: {emotion}
Assessed Risk Level: {risk}
Required Clinical Strategy to Teach the User:
{strategy}
Draft your response to the user's new message now:
"""
# 5. Groq Generation (via the LangChain LCEL chain)
final_response = self.generation_chain.invoke({
"augmented_prompt": augmented_prompt
})
return {
"session_id": session_id,
"response": final_response,
"emotion": emotion,
"confidence": prediction["confidence"],
"risk_level": risk,
}
# ─────────────────────────────────────────────────────────────
# RESOURCE CACHING
# ─────────────────────────────────────────────────────────────
@st.cache_resource
def get_safeguard_bot():
return SafeguardChatbot()
@st.cache_resource
def get_shap_explainer():
return MindGuardSHAPExplainer()
# ─────────────────────────────────────────────────────────────
# UI HELPERS (Badges & Formatting)
# ─────────────────────────────────────────────────────────────
RISK_COLORS = {
"High": ("#ff4b4b", "white"), "Medium": ("#ffa500", "white"), "Low": ("#21c354", "white")
}
EMOTION_COLORS = {
"Suicidal": "#ff4b4b", "Depression": "#e05260", "Anxiety": "#e07052", "Bipolar": "#c0392b",
"Stress": "#e67e22", "Personality disorder": "#9b59b6", "joy": "#21c354", "love": "#2ecc71",
"gratitude": "#27ae60", "optimism": "#16a085", "Normal": "#3498db", "neutral": "#3498db",
"sadness": "#e08000", "fear": "#e74c3c", "anger": "#c0392b"
}
def _emotion_badge(emotion: str) -> str:
color = EMOTION_COLORS.get(emotion, "#555555")
return f'🧠 {emotion}'
def _risk_badge(risk: str) -> str:
bg, fg = RISK_COLORS.get(risk, ("#888888", "white"))
icon = {"High": "🚨", "Medium": "⚠️", "Low": "✅"}.get(risk, "•")
return f'{icon} Risk: {risk}'
# ─────────────────────────────────────────────────────────────
# PAGE 1: CHAT COMPANION
# ─────────────────────────────────────────────────────────────
def render_chat():
st.title("🧠 Mindguard AI Companion")
st.markdown("Your clinical-grade, empathetic AI. Type a message or upload a voice note.")
bot = get_safeguard_bot()
if "messages" not in st.session_state:
st.session_state.messages = []
# One unique ID per browser session. Streamlit reruns the script often,
# but session_state survives those reruns.
# Replace this with the authenticated user ID when auth is added.
if "session_id" not in st.session_state:
st.session_state.session_id = f"streamlit_{uuid.uuid4().hex}"
# SIDEBAR: Audio & Controls
with st.sidebar:
st.header("🎙️ Voice Input")
audio_value = st.audio_input("Record a voice note")
if audio_value:
current_audio_size = len(audio_value.getvalue())
if ("last_audio_size" not in st.session_state or st.session_state.last_audio_size != current_audio_size):
st.session_state.last_audio_size = current_audio_size
os.makedirs(os.path.join(project_root, "data", "raw"), exist_ok=True)
temp_path = os.path.join(project_root, "data", "raw", "temp_streamlit.wav")
with open(temp_path, "wb") as f:
f.write(audio_value.getbuffer())
with st.spinner("Transcribing via Whisper…"):
transcribed_text = bot.audio_processor.transcribe(temp_path)
# --- DEFERRED EXECUTION FLAGS ---
st.session_state.latest_analyzable_text = transcribed_text
st.session_state.shap_is_stale = True
result = bot.generate_response(
transcribed_text,
st.session_state.messages,
session_id=st.session_state.session_id,
)
st.session_state.messages.append({"role": "user", "content": f"🎤 **Voice Note:** *{transcribed_text}*"})
st.session_state.messages.append({
"role": "assistant",
"content": result["response"],
"emotion": result["emotion"],
"risk_level": result["risk_level"],
"confidence": result["confidence"],
"session_id": result["session_id"],
})
st.divider()
if st.button("🗑️ Clear Chat"):
st.session_state.messages = []
st.session_state.pop("last_audio_size", None)
st.session_state.pop("latest_analyzable_text", None)
st.session_state.pop("shap_is_stale", None)
st.rerun()
# MAIN WINDOW: Chat History
for msg in st.session_state.messages:
with st.chat_message(msg["role"]):
st.markdown(msg["content"])
if msg["role"] == "assistant" and "emotion" in msg:
col1, col2 = st.columns([1, 1])
with col1: st.markdown(_emotion_badge(msg["emotion"]), unsafe_allow_html=True)
with col2: st.markdown(_risk_badge(msg.get("risk_level", "Unknown")), unsafe_allow_html=True)
# MAIN WINDOW: Text Input
if prompt := st.chat_input("How are you feeling right now?"):
st.session_state.messages.append({"role": "user", "content": prompt})
# --- DEFERRED EXECUTION FLAGS ---
st.session_state.latest_analyzable_text = prompt
st.session_state.shap_is_stale = True
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
with st.spinner("Diagnosing and retrieving clinical strategy…"):
result = bot.generate_response(
prompt,
st.session_state.messages,
session_id=st.session_state.session_id,
)
response = result["response"]
emotion = result["emotion"]
risk = result["risk_level"]
st.markdown(response)
col1, col2 = st.columns([1, 1])
with col1: st.markdown(_emotion_badge(emotion), unsafe_allow_html=True)
with col2: st.markdown(_risk_badge(risk), unsafe_allow_html=True)
st.session_state.messages.append({
"role": "assistant",
"content": response,
"emotion": emotion,
"risk_level": risk,
"confidence": result["confidence"],
"session_id": result["session_id"],
})
# ─────────────────────────────────────────────────────────────
# PAGE 2: CLINICAL DASHBOARD
# ─────────────────────────────────────────────────────────────
def render_dashboard():
st.title("📊 Clinical Overview")
st.markdown("Real-time emotional tracking and XAI reports for the current session.")
# NOTE: st.tabs() executes the code of EVERY tab on each rerun (even hidden
# ones), which would trigger SHAP as soon as the dashboard opens. A radio
# selector only runs the branch you actually picked.
view = st.radio(
"Dashboard view",
["📈 Analytics", "🔬 XAI Report"],
horizontal=True,
label_visibility="collapsed",
key="dashboard_view",
)
SHAP_HTML_PATH = os.path.join(project_root, "artifacts", "shap_report.html")
if view == "📈 Analytics":
history_data = []
for msg in st.session_state.get("messages", []):
if msg["role"] == "assistant" and "emotion" in msg:
history_data.append({
"diagnosed_emotion": msg["emotion"],
"risk_level": msg["risk_level"],
"timestamp": pd.Timestamp.now().strftime("%H:%M:%S")
})
if not history_data:
st.info("No data yet. Start chatting to generate live analytics.")
return
df = pd.DataFrame(history_data)
total = len(df)
high_risk = len(df[df["risk_level"] == "High"])
top_emotion = df["diagnosed_emotion"].mode()[0] if not df.empty else "N/A"
col1, col2, col3 = st.columns(3)
col1.metric("Interactions This Session", total)
col2.metric("High Risk Flags", high_risk, delta_color="inverse")
col3.metric("Primary Emotion", top_emotion)
st.divider()
st.subheader("Risk Level Summary")
risk_cols = st.columns(3)
for i, level in enumerate(["High", "Medium", "Low"]):
count = len(df[df["risk_level"] == level])
icons = {"High": "🚨", "Medium": "⚠️", "Low": "✅"}
risk_cols[i].metric(f"{icons[level]} {level}", count)
st.divider()
chart_col1, chart_col2 = st.columns(2)
with chart_col1:
st.subheader("Emotion Frequency")
st.bar_chart(df["diagnosed_emotion"].value_counts())
with chart_col2:
st.subheader("Risk Level Distribution")
st.bar_chart(df["risk_level"].value_counts())
else: # "🔬 XAI Report" — SHAP only runs when this view is selected
st.subheader("🔬 Last SHAP Word-Level Explanation")
# --- THE FIX: Lazy Loading Execution ---
if st.session_state.get("shap_is_stale", False) and "latest_analyzable_text" in st.session_state:
with st.spinner("Calculating XAI feature importance for the latest message... (This takes a moment)"):
shap_ex = get_shap_explainer()
shap_ex.generate_visual_report(st.session_state.latest_analyzable_text)
st.session_state.shap_is_stale = False
# Render the HTML if it exists
if os.path.exists(SHAP_HTML_PATH):
with open(SHAP_HTML_PATH, "r", encoding="utf-8") as f:
shap_html = f.read()
safe_html = f'{shap_html}
'
components.html(safe_html, height=500, scrolling=True)
st.caption(f"Report path: `{SHAP_HTML_PATH}`")
else:
st.info("Send a message in the Chat Companion tab to generate the first XAI report.")
# ─────────────────────────────────────────────────────────────
# MAIN ROUTER
# ─────────────────────────────────────────────────────────────
st.sidebar.title("🛡️ Mindguard AI")
st.sidebar.markdown("Welcome to the control panel.")
page = st.sidebar.radio("Navigation", ["💬 Chat Companion", "📊 Clinical Dashboard"])
if page == "💬 Chat Companion":
render_chat()
elif page == "📊 Clinical Dashboard":
render_dashboard()