Update app.py

app.py

@@ -4,7 +4,8 @@ import json
 import logging
 import re
 from datetime import datetime
-from typing import List, Dict, Optional
+from typing import List, Dict, Optional, Tuple
+from enum import Enum
 
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
@@ -20,7 +21,7 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger("TxAgentAPI")
 
 # App
-app = FastAPI(title="TxAgent API", version="2.1.0")
+app = FastAPI(title="TxAgent API", version="2.2.0")  # Version bump for new features
 
 app.add_middleware(
     CORSMiddleware,
@@ -36,10 +37,19 @@ class ChatRequest(BaseModel):
     history: Optional[List[Dict]] = None
     format: Optional[str] = "clean"
 
+# Enums
+class RiskLevel(str, Enum):
+    NONE = "none"
+    LOW = "low"
+    MODERATE = "moderate"
+    HIGH = "high"
+    SEVERE = "severe"
+
 # Globals
 agent = None
 patients_collection = None
 analysis_collection = None
+alerts_collection = None
 
 # Helpers
 def clean_text_response(text: str) -> str:
@@ -59,25 +69,21 @@ def extract_section(text: str, heading: str) -> str:
 def structure_medical_response(text: str) -> Dict:
     """Improved version that handles both markdown and plain text formats"""
     def extract_improved(text: str, heading: str) -> str:
-        # Try multiple patterns to match different heading formats
         patterns = [
-            rf"{re.escape(heading)}:\s*\n(.*?)(?=\n\s*\n|\Z)",  # Heading followed by content until double newline
-            rf"\*\*{re.escape(heading)}\*\*:\s*\n(.*?)(?=\n\s*\n|\Z)",  # Markdown bold heading
-            rf"{re.escape(heading)}[\s\-]+(.*?)(?=\n\s*\n|\Z)",  # Heading with dashes
-            rf"\n{re.escape(heading)}\s*\n(.*?)(?=\n\s*\n|\Z)"  # Heading on its own line
+            rf"{re.escape(heading)}:\s*\n(.*?)(?=\n\s*\n|\Z)",
+            rf"\*\*{re.escape(heading)}\*\*:\s*\n(.*?)(?=\n\s*\n|\Z)",
+            rf"{re.escape(heading)}[\s\-]+(.*?)(?=\n\s*\n|\Z)",
+            rf"\n{re.escape(heading)}\s*\n(.*?)(?=\n\s*\n|\Z)"
         ]
         
         for pattern in patterns:
             match = re.search(pattern, text, re.DOTALL | re.IGNORECASE)
             if match:
                 content = match.group(1).strip()
-                # Clean up any remaining markdown or special characters
                 content = re.sub(r'^\s*[\-\*]\s*', '', content, flags=re.MULTILINE)
                 return content
-        
         return ""
 
-    # Normalize the text first
     text = text.replace('**', '').replace('__', '')
     
     return {
@@ -90,6 +96,72 @@ def structure_medical_response(text: str) -> Dict:
         "recommendations": extract_improved(text, "Suggest Next Clinical Steps") or 
                          extract_improved(text, "Suggested Clinical Actions")
     }
+
+def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
+    """Analyze text for suicide risk factors and return assessment"""
+    suicide_keywords = [
+        'suicide', 'suicidal', 'kill myself', 'end my life', 
+        'want to die', 'self-harm', 'self harm', 'hopeless',
+        'no reason to live', 'plan to die'
+    ]
+    
+    # Check for explicit mentions
+    explicit_mentions = [kw for kw in suicide_keywords if kw in text.lower()]
+    
+    if not explicit_mentions:
+        return RiskLevel.NONE, 0.0, []
+    
+    # If found, ask AI for detailed assessment
+    assessment_prompt = (
+        "Assess the suicide risk level based on this text. "
+        "Consider frequency, specificity, and severity of statements. "
+        "Respond with JSON format: {\"risk_level\": \"low/moderate/high/severe\", "
+        "\"risk_score\": 0-1, \"factors\": [\"list of risk factors\"]}\n\n"
+        f"Text to assess:\n{text}"
+    )
+    
+    try:
+        response = agent.chat(
+            message=assessment_prompt,
+            history=[],
+            temperature=0.2,  # Lower temp for more deterministic responses
+            max_new_tokens=256
+        )
+        
+        # Extract JSON from response
+        json_match = re.search(r'\{.*\}', response, re.DOTALL)
+        if json_match:
+            assessment = json.loads(json_match.group())
+            return (
+                RiskLevel(assessment.get("risk_level", "none").lower()),
+                float(assessment.get("risk_score", 0)),
+                assessment.get("factors", [])
+            )
+    except Exception as e:
+        logger.error(f"Error in suicide risk assessment: {e}")
+    
+    # Fallback if JSON parsing fails
+    risk_score = min(0.1 * len(explicit_mentions), 0.9)  # Cap at 0.9 for fallback
+    if risk_score > 0.7:
+        return RiskLevel.HIGH, risk_score, explicit_mentions
+    elif risk_score > 0.4:
+        return RiskLevel.MODERATE, risk_score, explicit_mentions
+    return RiskLevel.LOW, risk_score, explicit_mentions
+
+async def create_alert(patient_id: str, risk_data: dict):
+    """Create an alert document in the database"""
+    alert_doc = {
+        "patient_id": patient_id,
+        "type": "suicide_risk",
+        "level": risk_data["level"],
+        "score": risk_data["score"],
+        "factors": risk_data["factors"],
+        "timestamp": datetime.utcnow(),
+        "acknowledged": False
+    }
+    await alerts_collection.insert_one(alert_doc)
+    logger.warning(f"⚠️ Created suicide risk alert for patient {patient_id}")
+
 def serialize_patient(patient: dict) -> dict:
     patient_copy = patient.copy()
     if "_id" in patient_copy:
@@ -101,13 +173,13 @@ async def analyze_patient(patient: dict):
         serialized = serialize_patient(patient)
         doc = json.dumps(serialized, indent=2)
         logger.info(f"🧾 Analyzing patient: {serialized.get('fhir_id')}")
-        logger.debug(f"🧠 Data passed to TxAgent:\n{doc[:1000]}")
 
+        # Main clinical analysis
         message = (
             "You are a clinical decision support AI.\n\n"
             "Given the patient document below:\n"
             "1. Summarize the patient's medical history.\n"
-            "2. Identify risks or red flags.\n"
+            "2. Identify risks or red flags (including mental health and suicide risk).\n"
             "3. Highlight missed diagnoses or treatments.\n"
             "4. Suggest next clinical steps.\n"
             f"\nPatient Document:\n{'-'*40}\n{doc[:10000]}"
@@ -115,19 +187,36 @@ async def analyze_patient(patient: dict):
 
         raw = agent.chat(message=message, history=[], temperature=0.7, max_new_tokens=1024)
         structured = structure_medical_response(raw)
-
+        
+        # Suicide risk assessment
+        risk_level, risk_score, risk_factors = detect_suicide_risk(raw)
+        suicide_risk = {
+            "level": risk_level.value,
+            "score": risk_score,
+            "factors": risk_factors
+        }
+        
+        # Store analysis
         analysis_doc = {
             "patient_id": serialized.get("fhir_id"),
             "timestamp": datetime.utcnow(),
             "summary": structured,
+            "suicide_risk": suicide_risk,
             "raw": raw
         }
+        
         await analysis_collection.update_one(
             {"patient_id": serialized.get("fhir_id")},
             {"$set": analysis_doc},
             upsert=True
         )
+        
+        # Create alert if risk is above threshold
+        if risk_level in [RiskLevel.MODERATE, RiskLevel.HIGH, RiskLevel.SEVERE]:
+            await create_alert(serialized.get("fhir_id"), suicide_risk)
+            
         logger.info(f"✅ Stored analysis for patient {serialized.get('fhir_id')}")
+
     except Exception as e:
         logger.error(f"Error analyzing patient: {e}")
 
@@ -139,7 +228,7 @@ async def analyze_all_patients():
 
 @app.on_event("startup")
 async def startup_event():
-    global agent, patients_collection, analysis_collection
+    global agent, patients_collection, analysis_collection, alerts_collection
 
     agent = TxAgent(
         model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
@@ -160,6 +249,7 @@ async def startup_event():
     db = get_mongo_client()["cps_db"]
     patients_collection = db["patients"]
     analysis_collection = db["patient_analysis_results"]
+    alerts_collection = db["clinical_alerts"]  # New collection for alerts
     logger.info("📡 Connected to MongoDB")
 
     asyncio.create_task(analyze_all_patients())
@@ -168,7 +258,8 @@ async def startup_event():
 async def status():
     return {
         "status": "running",
-        "timestamp": datetime.utcnow().isoformat()
+        "timestamp": datetime.utcnow().isoformat(),
+        "version": "2.2.0"
     }
 
 @app.post("/chat-stream")
@@ -201,4 +292,4 @@ async def chat_stream_endpoint(request: ChatRequest):
             logger.error(f"Streaming error: {e}")
             yield f"⚠️ Error: {e}"
 
-    return StreamingResponse(token_stream(), media_type="text/plain")
+    return StreamingResponse(token_stream(), media_type="text/plain")