Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import pipeline
 from huggingface_hub import login
 from openai import OpenAI
 import os
-import re
+import json
 
 # --- 1. SETUP ---
 hf_token = os.getenv("HF_TOKEN")
@@ -33,16 +33,17 @@ def detect_depression_risk(text: str) -> dict:
     """Analyzes text using Mental-Longformer (eRisk 2025)."""
     
     # --- THESIS LOGIC: AGGREGATION ---
-    # Strictly preserving your requirement for double newlines
+    # 1. We replace single newlines with double newlines for the MODEL
     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. Focus on hobbies/events."
+        desc = "Matches 'Isolated Control' group. High lexical diversity, focus on hobbies/events."
         color = "#10b981" # Green
     elif 0.40 <= prob < 0.60: 
         level = "Moderate Risk"
@@ -64,14 +65,7 @@ def detect_depression_risk(text: str) -> dict:
         "word_count": len(processed_text.split())
     }
 
-# --- 4. AGENT REASONING (With DeepSeek Cleaning) ---
-
-def clean_deepseek_output(text):
-    """Removes the <think> tags from DeepSeek R1 models."""
-    # Regex to remove <think>...</think> content (dotall to match newlines)
-    cleaned = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
-    return cleaned.strip()
-
+# --- 4. AGENT REASONING (SambaNova + Nebius Kimi) ---
 def agent_reasoning(text, risk_data, provider="SambaNova"):
     """
     Uses Sponsor APIs to generate the analysis report.
@@ -79,39 +73,43 @@ def agent_reasoning(text, risk_data, provider="SambaNova"):
     client = None
     model_id = None
     
-    # PROMPT
+    # SYSTEM PROMPT
     system_prompt = f"""
     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
+    CLINICAL DATA:
+    - Risk Level: {risk_data['risk_level']} ({risk_data['probability']:.1%})
+    - Biomarker: {risk_data['biomarker']}
     
     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 but professional. 
-    4. Keep it under 100 words.
+    3. Be compassionate. Max 100 words.
     """
 
     try:
-        # --- SPONSOR 1: NEBIUS (DeepSeek R1) ---
-        if provider == "Nebius (DeepSeek R1)":
+        # --- SPONSOR 1: NEBIUS (Kimi K2 Thinking) ---
+        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)
-            model_id = "deepseek-ai/DeepSeek-R1" # Standard R1 ID
+            client = OpenAI(
+                base_url="https://api.tokenfactory.nebius.com/v1/", 
+                api_key=api_key
+            )
+            model_id = "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)
+            client = OpenAI(
+                base_url="https://api.sambanova.ai/v1", 
+                api_key=api_key
+            )
             model_id = "Meta-Llama-3.3-70B-Instruct" 
 
         # EXECUTE
@@ -122,23 +120,16 @@ def agent_reasoning(text, risk_data, provider="SambaNova"):
                 {"role": "user", "content": "Analyze this."}
             ],
             temperature=0.6,
-            max_tokens=500 # Needs more tokens for thinking
+            max_tokens=300
         )
-        
-        raw_output = response.choices[0].message.content
-        
-        # CLEANUP: Remove the "Thinking" part if it's DeepSeek
-        final_output = clean_deepseek_output(raw_output)
-        
-        return final_output
+        return response.choices[0].message.content
 
     except Exception as e:
         return f"Reasoning Error ({provider}): {str(e)}"
 
-# --- 5. PIPELINE (Visual Report) ---
-
+# --- 5. PIPELINE ---
 def full_analysis_pipeline(user_text, location, provider):
-    if not user_text.strip(): return "Please enter text.", "No data"
+    if not user_text.strip(): return "Please enter text."
     
     # 1. Run Tool
     risk_data = detect_depression_risk(user_text)
@@ -152,60 +143,54 @@ def full_analysis_pipeline(user_text, location, provider):
     # 4. Color Logic
     color = "green" if risk_data['probability'] < 0.4 else "orange" if risk_data['probability'] < 0.6 else "red"
     
-    # 5. Build Dashboard (HTML) - Fixed Spacing
-    html_dashboard = f"""
-    <div style="padding: 20px; border-radius: 12px; background-color: {data['color']}15; border: 1px solid {data['color']};">
-        <h2 style="color: {data['color']}; margin: 0;">{data['risk_level']}</h2>
-        <h3 style="margin: 0;">{data['probability']*100:.1f}% Match</h3>
-        
-        <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: 15px; white-space: pre-wrap;">
-            <p><strong>🧠 Thesis Biomarker:</strong> {data['biomarker']}</p>
-            <p style="font-size: 0.9em; opacity: 0.8;">{data['description']}</p>
-        </div>
+    # 5. Build HTML/Markdown Report
+    # FIX: Added 'white-space: pre-wrap' to preserve newlines in display
+    return f"""
+    <div style="border-left: 6px solid {color}; padding-left: 15px; background-color: {color}10; border-radius: 5px; padding: 20px;">
+        <h2 style="color:{color}; margin:0;">{risk_data['risk_level']}</h2>
+        <h3 style="margin-top:5px;">Probability: {risk_data['probability']*100:.1f}%</h3>
     </div>
-    """
     
-    # 6. Build Text Report (Markdown)
-    report_markdown = f"""
-    ### 🤖 Agent Analysis ({provider})
+    ### 🧠 Agent Analysis ({provider})
+    <div style="white-space: pre-wrap; font-family: sans-serif;">
     {explanation}
+    </div>
     
     ---
-    ### 🆘 Recommended Resources ({location})
+    ### 🆘 Recommended Resources
     {resources}
     """
-    
-    return html_dashboard, report_markdown
 
 # --- 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."""
 
 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 ---
 with gr.Blocks(title="Depression Risk Agent") as demo:
     gr.Markdown("# 🧠 Early Depression Detection Agent (MCP)")
-    gr.Markdown("Agentic system using **Mental-Longformer** (Tool) + **Multi-Provider Reasoning**.")
-    gr.Markdown("⚡ **Powered by:** [SambaNova](https://sambanova.ai/) (Llama 3.3) & [Nebius](https://nebius.com/) (DeepSeek R1)")
+    gr.Markdown("Agentic system using **Mental-Longformer** (Tool) + **Multi-Provider Reasoning** (Nebius/SambaNova).")
     
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 1. User History")
             input_text = gr.Textbox(
                 label="User Timeline", 
-                lines=8, 
+                lines=10, 
                 placeholder="[Post 1] ...\n\n[Post 2] ...\n\n(Context >100 words recommended)"
             )
             
@@ -219,11 +204,14 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
             with gr.Row():
                 # RENAMED to be clearer
                 loc_dropdown = gr.Dropdown(["Global", "US", "Malaysia"], value="Malaysia", label="Crisis Resource Region")
-                # UPDATED PROVIDER LABELS
-                provider_dropdown = gr.Dropdown(["SambaNova", "Nebius (DeepSeek R1)"], value="SambaNova", label="Reasoning Brain")
+                provider_dropdown = gr.Dropdown(["SambaNova", "Nebius (Kimi K2)"], value="SambaNova", label="Reasoning Brain")
             
             submit = gr.Button("🚀 Run Analysis Agent", variant="primary", size="lg")
             
+        with gr.Column(scale=1):
+            gr.Markdown("### 4. Agent Response")
+            output = gr.Markdown(label="Response")
+            
             with gr.Accordion("🔧 MCP Tools Exposed", open=True):
                 gr.Markdown("""
                 The following functions are exposed to the MCP Client:
@@ -231,14 +219,7 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
                 - `get_crisis_resources(location)`: Returns localized help.
                 """)
 
-        with gr.Column(scale=1):
-            # OUTPUTS: ONE FOR HTML, ONE FOR MARKDOWN
-            out_dashboard = gr.HTML(label="Clinical Dashboard")
-            out_report = gr.Markdown(label="Agent Report")
-
-    # WIRING
-    # Ensure correct mapping: input_text -> full_analysis -> [dashboard, report]
-    submit.click(full_analysis_pipeline, inputs=[input_text, loc_dropdown, provider_dropdown], outputs=[out_dashboard, out_report])
+    submit.click(full_analysis_pipeline, inputs=[input_text, loc_dropdown, provider_dropdown], outputs=output)
     
     btn_low.click(lambda: example_low, None, input_text)
     btn_mod.click(lambda: example_mod, None, input_text)