Update app.py

app.py

@@ -9,7 +9,7 @@ hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-# --- 2. LOAD TOOL ---
+# --- 2. TOOL LOADING ---
 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)
@@ -23,6 +23,7 @@ def get_crisis_resources(location: str = "Global") -> str:
         "Malaysia": "🇲🇾 **Malaysia:** Befrienders KL: 03-76272929 | Talian Kasih: 15999",
         "Global": "🌍 **International:** [befrienders.org](https://www.befrienders.org)"
     }
+    # Fallback logic
     for key in resources:
         if location and key.lower() in location.lower():
             return resources[key]
@@ -30,238 +31,159 @@ 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 (CRITICAL) ---
+    # Thesis Aggregation Logic: 
     # We replace single newlines with double newlines so the model sees distinct posts
-    # This matches your Master's Thesis data preparation method.
-    processed_text = text.replace("\n", "\n\n") 
+    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 (Figure 4.15)
-    if prob < 0.40: 
-        level = "Low Risk"
-        biomarker = "Healthy External Focus"
-        desc = "Matches 'Isolated Control' group. High lexical diversity, focus on hobbies/events."
-        color = "#10b981" # Green
-    elif 0.40 <= prob < 0.60: 
-        level = "Moderate Risk"
-        biomarker = "Echo Chamber Interaction"
-        desc = "Matches 'Interactive Non-Depressed' group. Engaging in support forums but likely not clinically depressed (Supportive Responder)."
-        color = "#f59e0b" # Yellow
-    else: 
-        level = "High Risk"
-        biomarker = "Nocturnal & High-Effort"
-        desc = "Matches 'Depressed' cohort. Indicators: Nocturnal posting spikes (00-05 UTC), high-effort/low-frequency posting."
-        color = "#ef4444" # Red
+    if prob < 0.40: level = "Low Risk (Isolated Pattern)"
+    elif 0.40 <= prob < 0.60: level = "Moderate Risk (Interactive Echo Chamber Pattern)"
+    else: level = "High Risk (Depressed Pattern)"
         
-    return {
-        "probability": prob, 
-        "risk_level": level, 
-        "biomarker": biomarker,
-        "description": desc,
-        "color": color,
-        "word_count": len(processed_text.split())
-    }
-
-# --- 4. AGENT LOGIC (Dual State Management) ---
-
-def generate_response(api_history, risk_context, provider):
-    """Generates response using Sponsor API."""
-    if not risk_context:
-        risk_context = {"risk_level": "Unknown", "probability": 0.0, "description": "No analysis run yet."}
+    return {"probability": prob, "risk_level": level, "word_count": len(processed_text.split())}
 
+# --- 4. AGENT REASONING (UPDATED MODELS) ---
+def agent_reasoning(text, risk_data, provider="SambaNova"):
+    """
+    Uses Sponsor APIs to interpret the clinical data.
+    """
     client = None
-    model_id = None
+    model = None
     
-    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"
-    else: 
-        client = OpenAI(base_url="https://api.tokenfactory.nebius.com/v1/", api_key=os.getenv("NEBIUS_API_KEY"))
-        model_id = "moonshotai/Kimi-K2-Thinking"
-
+    # SYSTEM PROMPT (Thesis-Informed)
     system_prompt = f"""
-    You are 'Dr. Longformer', a specialized Clinical AI Assistant based on Hassan's 2025 Thesis.
-    CURRENT USER CONTEXT:
-    - Analyzed Risk: {risk_context['risk_level']} ({risk_context['probability']:.1%})
-    - Detected Pattern: {risk_context['description']}
-    YOUR GOAL: Provide supportive, scientifically-grounded chat. Max 2 sentences.
+    You are 'Dr. Longformer', an empathetic Clinical AI Research Agent.
+    
+    CLINICAL DATA (Mental-Longformer):
+    - Risk Level: {risk_data['risk_level']}
+    - Probability: {risk_data['probability']:.1%}
+    - Input Length: {risk_data['word_count']} words
+    
+    USER TEXT SNIPPET: "{text[:800]}..."
+    
+    INSTRUCTIONS:
+    1. Acknowledge the user's situation based on the text.
+    2. Explain the risk level using thesis terms: 'Nocturnal Posting' (High), 'Supportive Responder' (Moderate), or 'Healthy External Focus' (Low).
+    3. Be compassionate. Max 100 words.
     """
 
-    messages = [{"role": "system", "content": system_prompt}]
-    messages.extend(api_history)
-
     try:
+        # --- SPONSOR 1: NEBIUS (Kimi K2) ---
+        if provider == "Nebius (Kimi K2)":
+            api_key = os.getenv("NEBIUS_API_KEY")
+            if not api_key: return "⚠️ Nebius API Key missing."
+            
+            client = OpenAI(base_url="https://api.tokenfactory.nebius.com/v1/", api_key=api_key)
+            # UPDATED MODEL ID
+            model = "moonshotai/Kimi-K2-Thinking" 
+
+        # --- SPONSOR 2: SAMBANOVA (Llama 3.3) ---
+        elif provider == "SambaNova":
+            api_key = os.getenv("SAMBANOVA_API_KEY")
+            if not api_key: return "⚠️ SambaNova API Key missing."
+            
+            client = OpenAI(base_url="https://api.sambanova.ai/v1", api_key=api_key)
+            # UPDATED MODEL ID
+            model = "Meta-Llama-3.3-70B-Instruct" 
+
+        # EXECUTE
         response = client.chat.completions.create(
-            model=model_id, 
-            messages=messages, 
-            temperature=0.7, 
+            model=model,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": "Analyze this."}
+            ],
+            temperature=0.6,
             max_tokens=300
         )
         return response.choices[0].message.content
-    except Exception as e:
-        return f"⚠️ Error with {provider}: {str(e)}"
 
-# --- 5. UI ORCHESTRATION ---
+    except Exception as e:
+        return f"Reasoning Error ({provider}): {str(e)}"
 
-def run_analysis(text, location, provider):
-    if not text.strip(): return None, [], [], None
+# --- 5. PIPELINE ---
+def full_analysis_pipeline(user_text, location, provider):
+    if not user_text.strip(): return "Please enter text.", "No data", "No data"
     
-    # 1. Run Tool
-    data = detect_depression_risk(text)
+    risk_data = detect_depression_risk(user_text)
     resources = get_crisis_resources(location)
+    explanation = agent_reasoning(user_text, risk_data, provider)
     
-    # 2. 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>
-        <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>
-        </div>
-    </div>
-    """
-    
-    # 3. Agent Greeting
-    # Format for API: List of Dicts
-    api_history = [{"role": "user", "content": "I just ran the analysis. Please explain my results."}]
-    ai_msg = generate_response(api_history, data, provider)
-    api_history.append({"role": "assistant", "content": ai_msg})
+    # Color logic
+    color = "green" if risk_data['probability'] < 0.4 else "orange" if risk_data['probability'] < 0.6 else "red"
     
-    # Format for UI: List of Tuples [(None, AI_Message)]
-    # CRITICAL FIX: This standard format works on ALL Gradio versions
-    ui_history = [(None, ai_msg)]
-    
-    return html_dashboard, ui_history, api_history, data
-
-def user_chat(user_message, ui_history, api_history, risk_data, provider):
-    if not user_message: return ui_history, api_history, ""
-    
-    # 1. Update API History
-    api_history.append({"role": "user", "content": user_message})
-    
-    # 2. Update UI History (User message only first)
-    ui_history.append([user_message, None])
-    
-    # 3. Generate Response
-    ai_msg = generate_response(api_history, risk_data, provider)
-    
-    # 4. Update API History
-    api_history.append({"role": "assistant", "content": ai_msg})
+    return f"""
+    <div style="border-left: 5px solid {color}; padding-left: 10px;">
+        <h2 style="color:{color}; margin:0;">{risk_data['risk_level']}</h2>
+        <h3 style="margin-top:0;">Probability: {risk_data['probability']*100:.1f}%</h3>
+    </div>
     
-    # 5. Update UI History (Add AI response)
-    ui_history[-1][1] = ai_msg
+    ### 🧠 Agent Analysis ({provider})
+    {explanation}
     
-    return ui_history, api_history, ""
+    ### 🆘 Recommended Resources
+    {resources}
+    """
 
-# --- 6. EXAMPLES (RESTORED) ---
+# --- 6. EXAMPLES ---
 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."""
+Went hiking with some friends last weekend which was nice. Good to get outside and move around."""
 
 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 ---
+# --- 7. UI ---
 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** (Reasoning).")
-    gr.Markdown("⚡ **Powered by:** [SambaNova](https://sambanova.ai/) (Llama 3.3) & [Nebius](https://nebius.com/) (Kimi K2)")
-    
-    # Internal State
-    risk_state = gr.State(None)
-    api_state = gr.State([]) # Stores [{"role":...}] for the LLM
+    gr.Markdown("Agentic system using **Mental-Longformer** (Tool) + **Multi-Provider Reasoning** (Nebius/SambaNova).")
     
     with gr.Row():
-        # LEFT: INPUT
         with gr.Column(scale=1):
+            gr.Markdown("### 1. User History")
             input_text = gr.Textbox(
-                label="User Timeline (Paste posts here)", 
-                lines=8, 
-                placeholder="[Post 1] ...\n\n[Post 2] ...",
-                value=example_high
+                label="User Timeline", 
+                lines=10, 
+                placeholder="[Post 1] ...\n\n[Post 2] ...\n\n(Context >100 words recommended)"
             )
             
-            # EXAMPLES ROW
-            gr.Markdown("### 🔍 Try Thesis Patterns")
+            gr.Markdown("### 2. Thesis Patterns (Click to Load)")
             with gr.Row():
-                btn_low = gr.Button("🟢 Low Risk", size="sm")
-                btn_mod = gr.Button("🟡 Moderate", size="sm")
-                btn_high = gr.Button("🔴 High Risk", size="sm")
+                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")
+            gr.Markdown("### 3. Agent Settings")
             with gr.Row():
-                loc_drop = gr.Dropdown(["Global", "US", "Malaysia"], value="Malaysia", label="Crisis Resource Region")
-                prov_drop = gr.Dropdown(["SambaNova", "Nebius"], value="SambaNova", label="Agent Brain")
+                # RENAMED to be clearer
+                loc_dropdown = gr.Dropdown(["Global", "US", "Malaysia"], value="Global", label="Crisis Resource Region")
+                provider_dropdown = gr.Dropdown(["SambaNova", "Nebius (Kimi K2)"], value="SambaNova", label="Reasoning Brain")
             
-            analyze_btn = gr.Button("🚀 Run Clinical Analysis", variant="primary", size="lg")
+            submit = gr.Button("🚀 Run Analysis Agent", variant="primary", size="lg")
             
-            with gr.Accordion("🔧 MCP Tools Exposed", open=False):
-                gr.Markdown("- `detect_depression_risk`\n- `get_crisis_resources`")
-
-        # RIGHT: DASHBOARD & CHAT
         with gr.Column(scale=1):
-            dashboard = gr.HTML(label="Clinical Dashboard")
-            
-            # FIXED: REMOVED type="messages" to prevent crash. Uses standard Tuples.
-            chatbot = gr.Chatbot(label="Agent Chat", height=400)
+            gr.Markdown("### 4. Agent Response")
+            output = gr.Markdown(label="Response")
             
-            msg_input = gr.Textbox(label="Chat with Agent", placeholder="Ask about your results or get advice...")
-            send_btn = gr.Button("Send Message")
+            with gr.Accordion("🔧 MCP Tools Exposed", open=True):
+                gr.Markdown("""
+                The following functions are exposed to the MCP Client:
+                - `detect_depression_risk(text)`: Returns probability & thesis classification.
+                - `get_crisis_resources(location)`: Returns localized help.
+                """)
 
-    # WIRING
+    submit.click(full_analysis_pipeline, inputs=[input_text, loc_dropdown, provider_dropdown], outputs=output)
     
-    # 1. Analyze Button -> Updates Dashboard, Chatbot, Risk State
-    analyze_btn.click(
-        run_analysis,
-        inputs=[input_text, loc_drop, prov_drop],
-        outputs=[dashboard, chatbot, api_state, risk_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
-    # Note: We pass 'chatbot' as both input and output (it holds the history in 'tuples' format)
-    send_btn.click(
-        user_chat,
-        inputs=[msg_input, chatbot, api_state, risk_state, prov_drop],
-        outputs=[chatbot, api_state, msg_input]
-    )
-    msg_input.submit(
-        user_chat,
-        inputs=[msg_input, chatbot, api_state, risk_state, prov_drop],
-        outputs=[chatbot, api_state, msg_input]
-    )
 
 if __name__ == "__main__":
-    demo.launch(
-        mcp_server=True, 
-        theme=gr.themes.Soft()
-    )
+    demo.launch(mcp_server=True, theme=gr.themes.Soft())