Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers import pipeline
 from huggingface_hub import login
 from openai import OpenAI
 import os
+import re
 
 # --- 1. SETUP ---
 hf_token = os.getenv("HF_TOKEN")
@@ -33,7 +34,7 @@ def detect_depression_risk(text: str) -> dict:
     
     # --- THESIS LOGIC: AGGREGATION ---
     # Strictly preserving your requirement for double newlines
-    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)
@@ -63,7 +64,14 @@ def detect_depression_risk(text: str) -> dict:
         "word_count": len(processed_text.split())
     }
 
-# --- 4. AGENT REASONING (Distinct Models) ---
+# --- 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()
+
 def agent_reasoning(text, risk_data, provider="SambaNova"):
     """
     Uses Sponsor APIs to generate the analysis report.
@@ -85,26 +93,25 @@ def agent_reasoning(text, risk_data, provider="SambaNova"):
     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.
+    3. Be compassionate but professional. 
+    4. Keep it under 100 words.
     """
 
     try:
-        # --- SPONSOR 1: NEBIUS (DeepSeek R1 - REASONING MODEL) ---
+        # --- SPONSOR 1: NEBIUS (DeepSeek R1) ---
         if provider == "Nebius (DeepSeek R1)":
             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)
-            # Using DeepSeek R1 Distill for reasoning capabilities
-            model_id = "deepseek-ai/DeepSeek-R1-0528"
+            model_id = "deepseek-ai/DeepSeek-R1" # Standard R1 ID
 
-        # --- SPONSOR 2: SAMBANOVA (Llama 3.3 - SPEED MODEL) ---
+        # --- 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)
-            # Using latest Llama 3.3
             model_id = "Meta-Llama-3.3-70B-Instruct" 
 
         # EXECUTE
@@ -115,16 +122,23 @@ def agent_reasoning(text, risk_data, provider="SambaNova"):
                 {"role": "user", "content": "Analyze this."}
             ],
             temperature=0.6,
-            max_tokens=300
+            max_tokens=500 # Needs more tokens for thinking
         )
-        return response.choices[0].message.content
+        
+        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
 
     except Exception as e:
         return f"Reasoning Error ({provider}): {str(e)}"
 
-# --- 5. PIPELINE (Stable Report UI) ---
+# --- 5. PIPELINE (Visual Report) ---
+
 def full_analysis_pipeline(user_text, location, provider):
-    if not user_text.strip(): return "Please enter text."
+    if not user_text.strip(): return "Please enter text.", "No data"
     
     # 1. Run Tool
     risk_data = detect_depression_risk(user_text)
@@ -138,20 +152,34 @@ 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 HTML/Markdown Report
-    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>
+    # 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>
     </div>
+    """
     
-    ### 🧠 Agent Analysis ({provider})
+    # 6. Build Text Report (Markdown)
+    report_markdown = f"""
+    ### 🤖 Agent Analysis ({provider})
     {explanation}
     
     ---
-    ### 🆘 Recommended Resources
+    ### 🆘 Recommended Resources ({location})
     {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.
@@ -160,7 +188,7 @@ I also spent some time woodworking. I'm building a walnut coffee table and the g
 
 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."""
+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.
@@ -177,7 +205,7 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
             gr.Markdown("### 1. User History")
             input_text = gr.Textbox(
                 label="User Timeline", 
-                lines=10, 
+                lines=8, 
                 placeholder="[Post 1] ...\n\n[Post 2] ...\n\n(Context >100 words recommended)"
             )
             
@@ -196,10 +224,6 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
             
             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:
@@ -207,7 +231,14 @@ with gr.Blocks(title="Depression Risk Agent") as demo:
                 - `get_crisis_resources(location)`: Returns localized help.
                 """)
 
-    submit.click(full_analysis_pipeline, inputs=[input_text, loc_dropdown, provider_dropdown], outputs=output)
+        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])
     
     btn_low.click(lambda: example_low, None, input_text)
     btn_mod.click(lambda: example_mod, None, input_text)