Spaces:
Sleeping
Sleeping
| # ============================================================ | |
| # 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 | |
| # ───────────────────────────────────────────────────────────── | |
| def get_safeguard_bot(): | |
| return SafeguardChatbot() | |
| 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'<span style="background:{color};color:white;padding:3px 10px;border-radius:12px;font-size:13px;font-weight:600;">🧠 {emotion}</span>' | |
| def _risk_badge(risk: str) -> str: | |
| bg, fg = RISK_COLORS.get(risk, ("#888888", "white")) | |
| icon = {"High": "🚨", "Medium": "⚠️", "Low": "✅"}.get(risk, "•") | |
| return f'<span style="background:{bg};color:{fg};padding:3px 10px;border-radius:12px;font-size:13px;font-weight:600;">{icon} Risk: {risk}</span>' | |
| # ───────────────────────────────────────────────────────────── | |
| # 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'<div style="background-color: white; padding: 20px; border-radius: 10px;">{shap_html}</div>' | |
| 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() |