Update analysis.py

analysis.py

@@ -1,13 +1,31 @@
 from typing import Optional, Tuple, List
+from enum import Enum
 from config import agent, patients_collection, analysis_collection, alerts_collection, logger
 from models import RiskLevel
-from utils import structure_medical_response, compute_file_content_hash, compute_patient_data_hash, serialize_patient
+from utils import (
+    structure_medical_response, 
+    compute_file_content_hash, 
+    compute_patient_data_hash, 
+    serialize_patient,
+    broadcast_notification
+)
 from datetime import datetime
 import asyncio
 import json
 import re
 
+class NotificationType(str, Enum):
+    RISK_ALERT = "risk_alert"
+    SYSTEM = "system"
+    MESSAGE = "message"
+
+class NotificationStatus(str, Enum):
+    UNREAD = "unread"
+    READ = "read"
+    ARCHIVED = "archived"
+
 async def create_alert(patient_id: str, risk_data: dict):
+    """Create a new risk alert with notification metadata"""
     alert_doc = {
         "patient_id": patient_id,
         "type": "suicide_risk",
@@ -16,7 +34,6 @@ async def create_alert(patient_id: str, risk_data: dict):
         "factors": risk_data["factors"],
         "timestamp": datetime.utcnow(),
         "acknowledged": False,
-        # Facebook-like notification fields
         "notification": {
             "type": NotificationType.RISK_ALERT,
             "status": NotificationStatus.UNREAD,
@@ -27,83 +44,125 @@ async def create_alert(patient_id: str, risk_data: dict):
             "priority": "high" if risk_data["level"] in ["high", "severe"] else "medium"
         }
     }
-    await alerts_collection.insert_one(alert_doc)
-    
-    # Trigger real-time notification
-    await broadcast_notification(alert_doc["notification"])
     
-    logger.warning(f"⚠️ Created suicide risk alert for patient {patient_id}")
-    return alert_doc
-async def analyze_patient_report(patient_id: Optional[str], report_content: str, file_type: str, file_content: bytes):
+    try:
+        await alerts_collection.insert_one(alert_doc)
+        await broadcast_notification(alert_doc["notification"])
+        logger.warning(f"⚠️ Created suicide risk alert for patient {patient_id}")
+        return alert_doc
+    except Exception as e:
+        logger.error(f"Failed to create alert for patient {patient_id}: {str(e)}")
+        raise
+
+async def analyze_patient_report(
+    patient_id: Optional[str], 
+    report_content: str, 
+    file_type: str, 
+    file_content: bytes
+):
+    """Analyze a patient report and create alerts for risks"""
     identifier = patient_id if patient_id else compute_file_content_hash(file_content)
     report_data = {"identifier": identifier, "content": report_content, "file_type": file_type}
     report_hash = compute_patient_data_hash(report_data)
     logger.info(f"🧾 Analyzing report for identifier: {identifier}")
 
-    existing_analysis = await analysis_collection.find_one({"identifier": identifier, "report_hash": report_hash})
+    # Check for existing analysis
+    existing_analysis = await analysis_collection.find_one(
+        {"identifier": identifier, "report_hash": report_hash}
+    )
     if existing_analysis:
         logger.info(f"✅ No changes in report data for {identifier}, skipping analysis")
         return existing_analysis
 
-    prompt = (
-        "You are a clinical decision support AI. Analyze the following patient report:\n"
-        "1. Summarize the patient's medical history.\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 Report ({file_type}):\n{'-'*40}\n{report_content[:10000]}"
-    )
+    try:
+        # Generate analysis
+        prompt = (
+            "You are a clinical decision support AI. Analyze the following patient report:\n"
+            "1. Summarize the patient's medical history.\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 Report ({file_type}):\n{'-'*40}\n{report_content[:10000]}"
+        )
 
-    raw_response = agent.chat(
-        message=prompt,
-        history=[],
-        temperature=0.7,
-        max_new_tokens=1024
-    )
-    structured_response = structure_medical_response(raw_response)
+        raw_response = agent.chat(
+            message=prompt,
+            history=[],
+            temperature=0.7,
+            max_new_tokens=1024
+        )
+        structured_response = structure_medical_response(raw_response)
 
-    risk_level, risk_score, risk_factors = detect_suicide_risk(raw_response)
-    suicide_risk = {
-        "level": risk_level.value,
-        "score": risk_score,
-        "factors": risk_factors
-    }
+        # Detect suicide risk
+        risk_level, risk_score, risk_factors = detect_suicide_risk(raw_response)
+        suicide_risk = {
+            "level": risk_level.value,
+            "score": risk_score,
+            "factors": risk_factors
+        }
 
-    analysis_doc = {
-        "identifier": identifier,
-        "patient_id": patient_id,
-        "timestamp": datetime.utcnow(),
-        "summary": structured_response,
-        "suicide_risk": suicide_risk,
-        "raw": raw_response,
-        "report_hash": report_hash,
-        "file_type": file_type
-    }
+        # Store analysis
+        analysis_doc = {
+            "identifier": identifier,
+            "patient_id": patient_id,
+            "timestamp": datetime.utcnow(),
+            "summary": structured_response,
+            "suicide_risk": suicide_risk,
+            "raw": raw_response,
+            "report_hash": report_hash,
+            "file_type": file_type
+        }
 
-    await analysis_collection.update_one(
-        {"identifier": identifier, "report_hash": report_hash},
-        {"$set": analysis_doc},
-        upsert=True
-    )
+        await analysis_collection.update_one(
+            {"identifier": identifier, "report_hash": report_hash},
+            {"$set": analysis_doc},
+            upsert=True
+        )
 
-    if patient_id and risk_level in [RiskLevel.MODERATE, RiskLevel.HIGH, RiskLevel.SEVERE]:
-        await create_alert(patient_id, suicide_risk)
+        # Create alert if risk detected
+        if patient_id and risk_level in [RiskLevel.MODERATE, RiskLevel.HIGH, RiskLevel.SEVERE]:
+            await create_alert(patient_id, suicide_risk)
+
+        logger.info(f"✅ Stored analysis for identifier {identifier}")
+        return analysis_doc
 
-    logger.info(f"✅ Stored analysis for identifier {identifier}")
-    return analysis_doc
+    except Exception as e:
+        logger.error(f"Error analyzing report for {identifier}: {str(e)}")
+        error_alert = {
+            "identifier": identifier,
+            "type": "system_error",
+            "level": "high",
+            "message": f"Report analysis failed: {str(e)}",
+            "timestamp": datetime.utcnow(),
+            "acknowledged": False,
+            "notification": {
+                "type": NotificationType.SYSTEM,
+                "status": NotificationStatus.UNREAD,
+                "title": "Report Analysis Error",
+                "message": f"Failed to analyze report for {'patient ' + patient_id if patient_id else 'unknown identifier'}",
+                "icon": "❌",
+                "action_url": "/system/errors",
+                "priority": "high"
+            }
+        }
+        await alerts_collection.insert_one(error_alert)
+        raise
 
 async def analyze_patient(patient: dict):
+    """Analyze complete patient record and create alerts for risks"""
     try:
         serialized = serialize_patient(patient)
         patient_id = serialized.get("fhir_id")
         patient_hash = compute_patient_data_hash(serialized)
         logger.info(f"🧾 Analyzing patient: {patient_id}")
 
+        # Check for existing analysis
         existing_analysis = await analysis_collection.find_one({"patient_id": patient_id})
         if existing_analysis and existing_analysis.get("data_hash") == patient_hash:
             logger.info(f"✅ No changes in patient data for {patient_id}, skipping analysis")
             return
 
+        # Generate analysis
         doc = json.dumps(serialized, indent=2)
         message = (
             "You are a clinical decision support AI.\n\n"
@@ -118,6 +177,7 @@ async def analyze_patient(patient: dict):
         raw = agent.chat(message=message, history=[], temperature=0.7, max_new_tokens=1024)
         structured = structure_medical_response(raw)
         
+        # Detect suicide risk
         risk_level, risk_score, risk_factors = detect_suicide_risk(raw)
         suicide_risk = {
             "level": risk_level.value,
@@ -125,6 +185,7 @@ async def analyze_patient(patient: dict):
             "factors": risk_factors
         }
         
+        # Store analysis
         analysis_doc = {
             "identifier": patient_id,
             "patient_id": patient_id,
@@ -141,15 +202,36 @@ async def analyze_patient(patient: dict):
             upsert=True
         )
         
+        # Create alert if risk detected
         if risk_level in [RiskLevel.MODERATE, RiskLevel.HIGH, RiskLevel.SEVERE]:
             await create_alert(patient_id, suicide_risk)
             
         logger.info(f"✅ Stored analysis for patient {patient_id}")
 
     except Exception as e:
-        logger.error(f"Error analyzing patient: {e}")
+        logger.error(f"Error analyzing patient: {str(e)}")
+        error_alert = {
+            "patient_id": patient_id if 'patient_id' in locals() else "unknown",
+            "type": "system_error",
+            "level": "high",
+            "message": f"Patient analysis failed: {str(e)}",
+            "timestamp": datetime.utcnow(),
+            "acknowledged": False,
+            "notification": {
+                "type": NotificationType.SYSTEM,
+                "status": NotificationStatus.UNREAD,
+                "title": "Analysis Error",
+                "message": f"Failed to analyze patient {patient_id if 'patient_id' in locals() else 'unknown'}",
+                "icon": "❌",
+                "action_url": "/system/errors",
+                "priority": "high"
+            }
+        }
+        await alerts_collection.insert_one(error_alert)
+        raise
 
 def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
+    """Detect suicide risk level from text analysis"""
     suicide_keywords = [
         'suicide', 'suicidal', 'kill myself', 'end my life', 
         'want to die', 'self-harm', 'self harm', 'hopeless',
@@ -159,21 +241,24 @@ def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
     if not explicit_mentions:
         return RiskLevel.NONE, 0.0, []
     
-    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:
+        # Get AI 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}"
+        )
+        
         response = agent.chat(
             message=assessment_prompt,
             history=[],
             temperature=0.2,
             max_new_tokens=256
         )
+        
+        # Parse response
         json_match = re.search(r'\{.*\}', response, re.DOTALL)
         if json_match:
             assessment = json.loads(json_match.group())
@@ -185,6 +270,7 @@ def detect_suicide_risk(text: str) -> Tuple[RiskLevel, float, List[str]]:
     except Exception as e:
         logger.error(f"Error in suicide risk assessment: {e}")
     
+    # Fallback heuristic if AI assessment fails
     risk_score = min(0.1 * len(explicit_mentions), 0.9)
     if risk_score > 0.7:
         return RiskLevel.HIGH, risk_score, explicit_mentions