Update app.py

app.py

@@ -4,24 +4,24 @@ from huggingface_hub import login
 from openai import OpenAI
 import os
 
-# --- 1. SETUP ---
+# --- 1. AUTHENTICATION & SETUP ---
 hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-# --- 2. LOAD TOOL ---
+# --- 2. LOAD THE SPECIALIZED TOOL ---
 print("Loading Mental-Longformer...")
 model_name = "avtak/erisk-longformer-depression-v1"
 classifier = pipeline("text-classification", model=model_name, truncation=True, max_length=4096, top_k=None)
 
-# --- 3. LOGIC & TOOLS ---
+# --- 3. MCP TOOLS ---
 
 def get_crisis_resources(location: str = "Global") -> str:
     """Returns mental health resources based on location."""
     resources = {
         "US": "🇺🇸 **US:** Crisis Text Line: 741741 | Suicide Lifeline: 988",
         "Malaysia": "🇲🇾 **Malaysia:** Befrienders KL: 03-76272929 | Talian Kasih: 15999",
-        "Global": "🌍 **Global:** [befrienders.org](https://www.befrienders.org)"
+        "Global": "🌍 **International:** [befrienders.org](https://www.befrienders.org)"
     }
     for key in resources:
         if location and key.lower() in location.lower():
@@ -30,12 +30,12 @@ def get_crisis_resources(location: str = "Global") -> str:
 
 def detect_depression_risk(text: str) -> dict:
     """Analyzes text using Mental-Longformer (eRisk 2025)."""
-    # Thesis Logic: Aggregation
+    # Thesis Logic: User-Level Aggregation simulation
     processed_text = text.replace("\n", "\n\n")
     results = classifier(processed_text)[0]
     prob = next((r['score'] for r in results if r['label'] == 'LABEL_1'), 0)
     
-    # Thesis Thresholds & Biomarkers
+    # Thesis Thresholds (Figure 4.15)
     if prob < 0.40: 
         level = "Low Risk"
         biomarker = "Healthy External Focus"
@@ -49,7 +49,7 @@ def detect_depression_risk(text: str) -> dict:
     else: 
         level = "High Risk"
         biomarker = "Nocturnal & High-Effort"
-        desc = "Matches 'Depressed' cohort. Indicators: Nocturnal posting spikes (00-05 UTC), high-effort/low-frequency posting, semantic isolation."
+        desc = "Matches 'Depressed' cohort. Indicators: Nocturnal posting spikes (00-05 UTC), high-effort/low-frequency posting."
         color = "#ef4444" # Red
         
     return {
@@ -63,13 +63,11 @@ def detect_depression_risk(text: str) -> dict:
 
 # --- 4. AGENT CHAT LOGIC ---
 
-def generate_chat_response(history, risk_context, provider):
+def generate_chat_response(history_state, risk_context, provider):
     """
-    Generates a chat response using the selected Sponsor API.
-    Context-aware: It knows the risk analysis from the previous step.
+    Generates response using Sponsor API.
     """
     if not risk_context:
-        # Fallback if chat starts without analysis
         risk_context = {"risk_level": "Unknown", "probability": 0.0, "description": "No analysis run yet."}
 
     # CONFIGURE CLIENT
@@ -78,12 +76,12 @@ def generate_chat_response(history, risk_context, provider):
     
     if provider == "SambaNova":
         client = OpenAI(base_url="https://api.sambanova.ai/v1", api_key=os.getenv("SAMBANOVA_API_KEY"))
-        model_id = "Meta-Llama-3.3-70B-Instruct"
+        model_id = "Meta-Llama-3.3-70B-Instruct" # Latest Model
     else: # Nebius
         client = OpenAI(base_url="https://api.tokenfactory.nebius.com/v1/", api_key=os.getenv("NEBIUS_API_KEY"))
-        model_id = "moonshotai/Kimi-K2-Thinking"
+        model_id = "moonshotai/Kimi-K2-Thinking" # Requested Model
 
-    # SYSTEM PROMPT (Injects Thesis Context)
+    # SYSTEM PROMPT (Thesis-Informed)
     system_prompt = f"""
     You are 'Dr. Longformer', a specialized Clinical AI Assistant based on Hassan's 2025 Thesis.
     
@@ -94,17 +92,15 @@ def generate_chat_response(history, risk_context, provider):
     YOUR GOAL:
     Provide supportive, scientifically-grounded chat. 
     - If High Risk: Focus on sleep hygiene ('Nocturnal' findings) and breaking isolation.
-    - If Moderate Risk: Discuss burnout and setting boundaries ('Supportive Responder' fatigue).
+    - If Moderate Risk: Discuss burnout and setting boundaries.
     - If Low Risk: Encourage their hobbies.
     
     DO NOT diagnose. Use phrases like "Our analysis suggests..." or "Thesis patterns indicate..."
-    Be concise and conversational.
+    Be concise (max 2 sentences per reply) and conversational.
     """
 
-    # Format history for OpenAI API
     messages = [{"role": "system", "content": system_prompt}]
-    for msg in history:
-        messages.append({"role": msg["role"], "content": msg["content"]})
+    messages.extend(history_state)
 
     try:
         response = client.chat.completions.create(
@@ -120,30 +116,26 @@ def generate_chat_response(history, risk_context, provider):
 # --- 5. UI ORCHESTRATION ---
 
 def run_analysis(text, location, provider):
-    if not text.strip(): return None, None, None
+    if not text.strip(): return None, None, None, None
     
     # 1. Run Tool
     data = detect_depression_risk(text)
     resources = get_crisis_resources(location)
     
-    # 2. Create Visual Dashboard (HTML)
+    # 2. Create Visual Dashboard
     html_dashboard = f"""
     <div style="padding: 20px; border-radius: 12px; background-color: {data['color']}15; border: 1px solid {data['color']};">
         <div style="display: flex; justify-content: space-between; align-items: center;">
             <h2 style="color: {data['color']}; margin: 0;">{data['risk_level']}</h2>
             <h3 style="margin: 0;">{data['probability']*100:.1f}% Match</h3>
         </div>
-        
-        <!-- Progress Bar -->
         <div style="background-color: #e5e7eb; border-radius: 9999px; height: 10px; margin-top: 10px; width: 100%;">
             <div style="background-color: {data['color']}; height: 10px; border-radius: 9999px; width: {data['probability']*100}%;"></div>
         </div>
-        
         <div style="margin-top: 20px;">
             <p><strong>🧠 Thesis Reflection:</strong> {data['biomarker']}</p>
             <p style="opacity: 0.8; font-size: 0.9em;">{data['description']}</p>
         </div>
-        
         <div style="margin-top: 15px; padding-top: 15px; border-top: 1px solid #ccc;">
             <p><strong>🆘 Local Resources ({location}):</strong></p>
             <pre style="font-family: inherit; white-space: pre-wrap; margin: 0;">{resources}</pre>
@@ -151,30 +143,47 @@ def run_analysis(text, location, provider):
     </div>
     """
     
-    # 3. Generate Initial Agent Greeting
-    greeting_prompt = [{"role": "user", "content": "I just ran the analysis. Please explain my results and start a conversation."}]
-    initial_response = generate_chat_response(greeting_prompt, data, provider)
+    # 3. Agent Greeting
+    history_state = []
+    initial_prompt = "I just ran the analysis. Please explain my results and start a conversation."
+    ai_response = generate_chat_response([{"role": "user", "content": initial_prompt}], data, provider)
     
-    # Return: Visual HTML, Chat History (with greeting), State Data
-    return html_dashboard, [{"role": "assistant", "content": initial_response}], data
+    history_state.append({"role": "assistant", "content": ai_response})
+    ui_history = [(None, ai_response)]
+    
+    return html_dashboard, ui_history, history_state, data
 
-def user_chat(user_message, history, risk_data, provider):
-    if not user_message: return history, ""
+def user_chat(user_message, ui_history, history_state, risk_data, provider):
+    if not user_message: return ui_history, history_state, ""
     
-    # Add user message
-    history.append({"role": "user", "content": user_message})
+    history_state.append({"role": "user", "content": user_message})
+    ui_history.append((user_message, None))
     
-    # Generate response
-    ai_msg = generate_chat_response(history, risk_data, provider)
+    ai_msg = generate_chat_response(history_state, risk_data, provider)
     
-    # Add AI message
-    history.append({"role": "assistant", "content": ai_msg})
+    history_state.append({"role": "assistant", "content": ai_msg})
+    ui_history[-1] = (user_message, ai_msg)
     
-    return history, ""
+    return ui_history, history_state, ""
+
+# --- 6. EXAMPLES (RESTORED) ---
+
+example_low = """The new update for the Linux kernel (6.8) finally addressed the driver issues I was seeing with Realtek cards. I read the changelog on kernel.org and noticed they patched the module specific to the rtl8821ce chipset.
+
+On a different note, the defensive stats for the basketball team have improved significantly. Allowing 15% fewer points per possession is a game changer.
+
+I also spent some time woodworking. I'm building a walnut coffee table and the grain is tricky to plane. I might switch to a polyurethane finish for durability."""
+
+example_mod = """Things have been pretty busy at work. We're in the middle of a big project and deadlines are tight. Staying a bit later than usual but that's just how it goes sometimes. Team is handling it well overall.
+
+Sleep could be better. Been staying up too late scrolling my phone. Need to work on that. Usually feel okay once I get moving in the morning though.
+
+Went hiking with some friends last weekend which was nice. Good to get outside and move around. We're talking about doing another trip next month. Weather should be better by then."""
 
-# --- 6. EXAMPLES ---
 example_high = """I don't know why I even bother getting out of bed anymore. I slept for 12 hours yesterday and I'm still exhausted. It feels like my limbs weigh a thousand pounds.
+
 I've been avoiding my friends for weeks. They keep texting me, but I can't bring myself to reply. The thought of socializing is terrifying.
+
 Everything feels like a shade of grey. I can't concentrate on my work. I feel like I'm drowning while everyone else is breathing fine. I was up until 4 AM again last night just staring at the ceiling."""
 
 # --- 7. UI LAYOUT ---
@@ -182,8 +191,8 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
     gr.Markdown("# 🧠 Early Depression Detection Agent (MCP)")
     gr.Markdown("Agentic system using **Mental-Longformer** (Tool) + **SambaNova/Nebius** (Chat).")
     
-    # State to hold the analysis data across chat turns
     risk_state = gr.State(None)
+    chat_state = gr.State([])
     
     with gr.Row():
         # LEFT: INPUT
@@ -195,6 +204,15 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
                 value=example_high
             )
             
+            # EXAMPLES ROW (Restored)
+            gr.Markdown("### 🔍 Try Thesis Patterns")
+            with gr.Row():
+                btn_low = gr.Button("🟢 Low Risk", size="sm", variant="secondary")
+                btn_mod = gr.Button("🟡 Moderate", size="sm", variant="secondary")
+                btn_high = gr.Button("🔴 High Risk", size="sm", variant="secondary")
+            
+            # SETTINGS ROW
+            gr.Markdown("### ⚙️ Settings")
             with gr.Row():
                 loc_drop = gr.Dropdown(["Global", "US", "Malaysia"], value="Malaysia", label="Region")
                 prov_drop = gr.Dropdown(["SambaNova", "Nebius"], value="SambaNova", label="Agent Brain")
@@ -206,34 +224,39 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
 
         # RIGHT: DASHBOARD & CHAT
         with gr.Column(scale=1):
-            # The "Reflection" Dashboard
             dashboard = gr.HTML(label="Clinical Dashboard")
-            
-            # The Chat Interface
-            chatbot = gr.Chatbot(type="messages", label="Agent Chat", height=400)
+            chatbot = gr.Chatbot(label="Agent Chat", height=400)
             msg_input = gr.Textbox(label="Chat with Agent", placeholder="Ask about your results or get advice...")
             send_btn = gr.Button("Send Message")
 
     # LOGIC WIRING
     
-    # 1. Analyze Button -> Updates Dashboard & Starts Chat
+    # 1. Analyze Button
     analyze_btn.click(
         run_analysis,
         inputs=[input_text, loc_drop, prov_drop],
-        outputs=[dashboard, chatbot, risk_state]
+        outputs=[dashboard, chatbot, chat_state, risk_state]
     )
     
-    # 2. Send Message -> Updates Chat using State
+    # 2. Example Buttons
+    btn_low.click(lambda: example_low, None, input_text)
+    btn_mod.click(lambda: example_mod, None, input_text)
+    btn_high.click(lambda: example_high, None, input_text)
+    
+    # 3. Chat Interactions
     send_btn.click(
         user_chat,
-        inputs=[msg_input, chatbot, risk_state, prov_drop],
-        outputs=[chatbot, msg_input]
+        inputs=[msg_input, chatbot, chat_state, risk_state, prov_drop],
+        outputs=[chatbot, chat_state, msg_input]
     )
     msg_input.submit(
         user_chat,
-        inputs=[msg_input, chatbot, risk_state, prov_drop],
-        outputs=[chatbot, msg_input]
+        inputs=[msg_input, chatbot, chat_state, risk_state, prov_drop],
+        outputs=[chatbot, chat_state, msg_input]
     )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True, theme=gr.themes.Soft())
+    demo.launch(
+        mcp_server=True, 
+        theme=gr.themes.Soft()
+    )