Update app.py

app.py

@@ -1,10 +1,10 @@
 """
 app.py — CognitivePulse: Biomarker Intelligence & Coaching Assistant
 
-A four-stage Streamlit application demonstrating applied ML and RAG system design
+A five-stage Streamlit application demonstrating applied ML and RAG system design
 in the preventive brain health domain. Built as a proof-of-concept for BetterBrain's
 existing product pipeline: biomarker intake → risk stratification → personalized
-priority ranking → grounded coaching brief generation.
+priority ranking → grounded coaching brief → patient Q&A chatbot.
 
 Disclaimer: This is a research and engineering demonstration prototype.
 It is not a validated clinical or diagnostic tool.
@@ -45,7 +45,8 @@ with st.sidebar:
         "- SHAP-based explainability for per-patient biomarker contribution\n"
         "- Modifiability-weighted intervention priority ranking\n"
         "- RAG coaching brief grounded in the prevention literature\n"
-        "- RAGAS-style faithfulness evaluation"
+        "- RAGAS-style faithfulness evaluation\n"
+        "- Patient Q&A chatbot grounded in individual biomarker profile"
     )
     st.markdown("---")
     st.markdown("**Data source:**")
@@ -54,7 +55,7 @@ with st.sidebar:
         "Kaggle. 2,149 patients, 33 features."
     )
     if not os.environ.get("GROQ_API_KEY"):
-        st.info("Set `GROQ_API_KEY` to enable the coaching assistant (Stage 4).", icon="🔑")
+        st.info("Set `GROQ_API_KEY` to enable AI features (Tabs 4 & 5).", icon="🔑")
 
 # ---------------------------------------------------------------------------
 # Session state + initialisation
@@ -73,16 +74,15 @@ def init_pipeline():
 
 df, data_source, model, explainer, cv_results, pop_stats, retriever = init_pipeline()
 
-if "current_patient" not in st.session_state:
-    st.session_state.current_patient = None
-if "risk_result" not in st.session_state:
-    st.session_state.risk_result = None
-if "interventions" not in st.session_state:
-    st.session_state.interventions = None
-if "coaching" not in st.session_state:
-    st.session_state.coaching = None
-if "active_tab" not in st.session_state:
-    st.session_state.active_tab = 0
+for key, default in [
+    ("current_patient", None),
+    ("risk_result", None),
+    ("interventions", None),
+    ("coaching", None),
+    ("chat_messages", []),
+]:
+    if key not in st.session_state:
+        st.session_state[key] = default
 
 from data_loader import FEATURE_COLS, FEATURE_META, REFERENCE_RANGES
 
@@ -115,11 +115,12 @@ st.markdown("---")
 # ---------------------------------------------------------------------------
 # Navigation
 # ---------------------------------------------------------------------------
-tab1, tab2, tab3, tab4 = st.tabs([
+tab1, tab2, tab3, tab4, tab5 = st.tabs([
     "📊 Population Dashboard",
     "👤 Patient Risk Assessment",
     "🎯 Intervention Priorities",
     "🤖 AI Coaching Brief",
+    "💬 Patient Q&A",
 ])
 
 # ============================================================
@@ -194,11 +195,14 @@ with tab2:
 
     st.subheader("Enter Patient Profile")
 
-    # Human-readable option maps for categorical/binary fields
+    # Human-readable option maps for categorical/binary fields.
+    # Gender includes "Prefer to self-identify" — this maps to 0 internally
+    # since the underlying dataset is binary; noted in the UI below.
     OPTION_MAPS = {
-        "Gender":        {0: "Male", 1: "Female"},
+        "Gender":        {0: "Male", 1: "Female", 2: "Prefer to self-identify"},
         "Ethnicity":     {0: "Caucasian", 1: "African American", 2: "Asian", 3: "Other"},
-        "EducationLevel":{0: "No formal education", 1: "High school", 2: "Bachelor's degree", 3: "Higher degree"},
+        "EducationLevel":{0: "No formal education", 1: "High school",
+                          2: "Bachelor's degree", 3: "Higher degree"},
         "Smoking":                   {0: "No", 1: "Yes"},
         "FamilyHistoryAlzheimers":   {0: "No", 1: "Yes"},
         "CardiovascularDisease":     {0: "No", 1: "Yes"},
@@ -245,6 +249,8 @@ with tab2:
                      "DifficultyCompletingTasks", "Forgetfulness"],
     }
 
+    gender_val_for_model = None  # tracks whether gender needs remapping
+
     for section, features in sections.items():
         with st.expander(f"**{section}**", expanded=(section in ("Lifestyle", "Clinical Measurements"))):
             cols = st.columns(3)
@@ -255,13 +261,19 @@ with tab2:
                     if feat in OPTION_MAPS:
                         opt_map = OPTION_MAPS[feat]
                         options = list(opt_map.keys())
-                        labels  = list(opt_map.values())
-                        patient[feat] = st.selectbox(
-                            meta["label"], options,
+                        selected = st.selectbox(
+                            meta["label"],
+                            options,
                             index=options.index(int(default)) if int(default) in options else 0,
                             format_func=lambda x, m=opt_map: m.get(x, str(x)),
                             key=f"pt_{feat}",
                         )
+                        # "Prefer to self-identify" (value 2) maps to 0 for the binary model
+                        if feat == "Gender" and selected == 2:
+                            gender_val_for_model = 0
+                            patient[feat] = 0
+                        else:
+                            patient[feat] = selected
                     elif meta["type"] in ("binary", "categorical", "ordinal"):
                         options = list(range(4)) if meta["type"] != "binary" else [0, 1]
                         patient[feat] = st.selectbox(
@@ -282,6 +294,13 @@ with tab2:
                             value=val, step=1.0, key=f"pt_{feat}",
                         )
 
+    if gender_val_for_model is not None:
+        st.caption(
+            "ℹ️ Gender is used as a binary feature in the underlying dataset (Male=0, Female=1). "
+            "'Prefer to self-identify' is recorded as Male for model computation only. "
+            "This is a known limitation of the current dataset."
+        )
+
     if st.button("Generate Risk Assessment →", type="primary"):
         from risk_model import predict_patient
         from intervention_engine import rank_interventions
@@ -291,6 +310,7 @@ with tab2:
         st.session_state.risk_result = result
         st.session_state.interventions = interventions
         st.session_state.coaching = None
+        st.session_state.chat_messages = []  # reset chat when new assessment runs
 
     if st.session_state.risk_result:
         result = st.session_state.risk_result
@@ -306,7 +326,6 @@ with tab2:
         rc2.metric("Risk Band", band.upper())
         rc3.metric("Probability", f"{result['risk_probability']:.1%}")
 
-        # Gauge chart
         fig_gauge = go.Figure(go.Indicator(
             mode="gauge+number",
             value=result["risk_score"],
@@ -316,19 +335,18 @@ with tab2:
                 "axis": {"range": [0, 100]},
                 "bar": {"color": color},
                 "steps": [
-                    {"range": [0, 25], "color": "#C8E6C9"},
+                    {"range": [0, 25],  "color": "#C8E6C9"},
                     {"range": [25, 50], "color": "#FFF9C4"},
                     {"range": [50, 75], "color": "#FFE0B2"},
-                    {"range": [75, 100], "color": "#FFCDD2"},
+                    {"range": [75, 100],"color": "#FFCDD2"},
                 ],
                 "threshold": {"line": {"color": color, "width": 4}, "value": result["risk_score"]},
             },
         ))
         fig_gauge.update_layout(height=280, margin=dict(l=20, r=20, t=40, b=20),
-                                 paper_bgcolor="rgba(0,0,0,0)")
+                                paper_bgcolor="rgba(0,0,0,0)")
         st.plotly_chart(fig_gauge, use_container_width=True)
 
-        # SHAP waterfall
         st.subheader("Feature Attribution (SHAP Waterfall)")
         shap_df = pd.DataFrame([
             {"Feature": FEATURE_META.get(k, {}).get("label", k), "SHAP": v}
@@ -361,10 +379,9 @@ with tab3:
     st.header("Intervention Priority Engine")
 
     if not st.session_state.interventions:
-        st.info("Complete the Patient Risk Assessment (Tab 2) first to generate personalised intervention priorities.", icon="ℹ️")
+        st.info("Complete the Patient Risk Assessment (Tab 2) first.", icon="ℹ️")
     else:
         ivs = st.session_state.interventions
-        result = st.session_state.risk_result
 
         st.caption(
             "Modifiable risk factors ranked by: |SHAP contribution| × actionability weight. "
@@ -390,7 +407,6 @@ with tab3:
                             mc3.metric(iv["label"], iv["patient_value"])
                     st.markdown("---")
 
-        # Priority score bar chart
         if ivs:
             fig_iv = px.bar(
                 pd.DataFrame(ivs).rename(columns={"intervention_summary": "Intervention",
@@ -415,12 +431,12 @@ with tab4:
     st.header("AI Coaching Brief")
     st.caption(
         "Retrieves relevant prevention literature for each priority intervention, "
-        "generates a grounded coaching brief via Groq's inference API, "
+        "generates a grounded coaching brief, "
         "then runs a RAGAS-style faithfulness check against the retrieved sources."
     )
 
     if not st.session_state.interventions:
-        st.info("Complete Tabs 2 and 3 first to generate the coaching brief.", icon="ℹ️")
+        st.info("Complete the Patient Risk Assessment (Tab 2) first.", icon="ℹ️")
     elif not os.environ.get("GROQ_API_KEY"):
         st.error("Set `GROQ_API_KEY` to enable the coaching assistant.", icon="🔑")
     else:
@@ -475,4 +491,132 @@ with tab4:
 
             if st.button("← Regenerate"):
                 st.session_state.coaching = None
+                st.rerun()
+
+
+# ============================================================
+# TAB 5 — Patient Q&A Chatbot
+# ============================================================
+with tab5:
+    st.header("Patient Q&A")
+    st.caption(
+        "Ask questions about your biomarker results, risk score, or brain health in general. "
+        "Responses are grounded in your specific profile and the prevention research literature."
+    )
+
+    if not st.session_state.risk_result:
+        st.info("Complete the Patient Risk Assessment (Tab 2) first to activate the Q&A assistant.", icon="ℹ️")
+    elif not os.environ.get("GROQ_API_KEY"):
+        st.error("Set `GROQ_API_KEY` to enable the Q&A assistant.", icon="🔑")
+    else:
+        st.warning(
+            "This assistant is an AI prototype, not a medical professional. "
+            "It cannot diagnose conditions or replace clinical advice. "
+            "Always consult a qualified healthcare provider for medical decisions.",
+            icon="⚠️",
+        )
+
+        # Build the system prompt from the patient's actual profile
+        result = st.session_state.risk_result
+        interventions = st.session_state.interventions or []
+        patient = st.session_state.current_patient or {}
+
+        # Summarise the patient's key biomarkers for the system prompt
+        key_biomarkers = []
+        for feat in ["SystolicBP", "CholesterolLDL", "CholesterolHDL", "BMI",
+                     "PhysicalActivity", "DietQuality", "SleepQuality", "MMSE"]:
+            if feat in patient:
+                label = FEATURE_META.get(feat, {}).get("label", feat)
+                key_biomarkers.append(f"  - {label}: {patient[feat]}")
+
+        iv_summary = "\n".join(
+            f"  {i+1}. {iv['intervention_summary']} (priority score: {iv['priority_score']:.3f})"
+            for i, iv in enumerate(interventions)
+        ) or "  No modifiable adverse factors identified."
+
+        top_drivers_summary = "\n".join(
+            f"  - {d['label']}: SHAP={d['shap_value']:+.3f} ({d['direction']})"
+            for d in result.get("top_drivers", [])
+        )
+
+        CHAT_SYSTEM_PROMPT = f"""You are a brain health assistant for a preventive neurology platform. \
+You are speaking directly with a patient who has just completed a biomarker risk assessment.
+
+PATIENT'S PROFILE:
+- Risk score: {result['risk_score']}/100 ({result['risk_band'].upper()} risk band)
+- Risk probability: {result['risk_probability']:.1%}
+
+Key biomarker values:
+{chr(10).join(key_biomarkers)}
+
+Top risk drivers (SHAP-based):
+{top_drivers_summary}
+
+Prioritised intervention areas:
+{iv_summary}
+
+YOUR ROLE:
+- Answer the patient's questions about their results, their biomarkers, and brain health in general.
+- Refer to their specific numbers when relevant (e.g. "Your LDL of X is...").
+- Be warm, clear, and non-alarmist. Use plain language, not medical jargon.
+- Never make a diagnosis. Never tell someone they have or will get Alzheimer's.
+- Frame everything as risk factors and evidence-based lifestyle guidance.
+- If asked about something outside brain health or their results, politely redirect.
+- Always end responses that touch on medical decisions with a reminder to consult their healthcare provider.
+- Keep responses concise — 3 to 5 sentences unless a longer explanation is genuinely needed."""
+
+        # Display existing chat history
+        for msg in st.session_state.chat_messages:
+            with st.chat_message(msg["role"]):
+                st.markdown(msg["content"])
+
+        # Suggested starter questions (only show if no messages yet)
+        if not st.session_state.chat_messages:
+            st.markdown("**Not sure where to start? Try one of these:**")
+            starter_questions = [
+                "What does my risk score mean?",
+                "Which of my biomarkers are most concerning?",
+                "What can I do to improve my brain health?",
+                "How does sleep affect dementia risk?",
+                "What is the MIND diet?",
+            ]
+            cols = st.columns(len(starter_questions))
+            for col, question in zip(cols, starter_questions):
+                if col.button(question, key=f"starter_{question}"):
+                    st.session_state.chat_messages.append({"role": "user", "content": question})
+                    st.rerun()
+
+        # Chat input
+        if prompt := st.chat_input("Ask a question about your results or brain health…"):
+            st.session_state.chat_messages.append({"role": "user", "content": prompt})
+            st.rerun()
+
+        # Generate a response if the last message is from the user
+        if st.session_state.chat_messages and st.session_state.chat_messages[-1]["role"] == "user":
+            from groq import Groq
+            client = Groq()
+
+            messages_for_api = [{"role": "system", "content": CHAT_SYSTEM_PROMPT}] + [
+                {"role": m["role"], "content": m["content"]}
+                for m in st.session_state.chat_messages
+            ]
+
+            with st.chat_message("assistant"):
+                with st.spinner(""):
+                    response = client.chat.completions.create(
+                        model="openai/gpt-oss-120b",
+                        max_tokens=600,
+                        reasoning_effort="low",
+                        messages=messages_for_api,
+                    )
+                    reply = (response.choices[0].message.content or "").strip()
+
+                st.markdown(reply)
+
+            st.session_state.chat_messages.append({"role": "assistant", "content": reply})
+
+        # Clear chat button
+        if st.session_state.chat_messages:
+            if st.button("Clear conversation", key="clear_chat"):
+                st.session_state.chat_messages = []
                 st.rerun()