Update app.py

app.py

@@ -48,15 +48,33 @@ class MedicalTriageBot:
         self.model = None
         self.embeddings = None
         self.vector_store = None
-        self.chunk_size = 300
-        self.chunk_overlap = 100
-        self.num_relevant_chunks = 3
+        self.chunk_size = 512
+        self.chunk_overlap = 128
+        self.num_relevant_chunks = 5
         self.last_interaction_time = time.time()
         self.interaction_cooldown = 1.0
         self.safety_phrases = [
             "999", "111", "emergency", "GP", "NHS",
             "consult a doctor", "seek medical attention"
         ]
+        self.triage_levels = {
+            'Emergency': [
+                'pediatric', 'stroke', 'cardiac', 'unconscious', 'suicidal',
+                'psychotic', 'seizure', 'overdose', 'cyanosis'
+            ],
+            'Urgent': [
+                'pregnancy', 'fracture', 'asthma', 'withdrawal',
+                'postpartum', 'harm thoughts', 'panic attack'
+            ],
+            'GP Care': [
+                'eczema', 'ocd', 'depression', 'anxiety', 'PTSD',
+                'bipolar', 'migraine', 'chronic'
+            ],
+            'Self-Care': [
+                'cold', 'rash', 'stress', 'mild', 'situational',
+                'acne', 'insomnia'
+            ]
+        }
         self.current_case = None
         self.location_services = {
             "London": {"A&E": ["St Thomas' Hospital", "Royal London Hospital"],
@@ -69,6 +87,45 @@ class MedicalTriageBot:
             ("emergency_signs", "Any difficulty breathing, chest pain, or confusion?")
         ]
 
+    def _determine_triage_level(self, context):
+        context_lower = context.lower()
+        for level, keywords in triage_levels.items():
+            if any(kw in context_lower for kw in keywords):
+                return level
+        return 'GP Care'
+
+    def _parse_medical_context(self, context):
+        components = {
+            'advice': [], 'red_flags': [], 
+            'timeframe': '24 hours', 'special_considerations': []
+        }
+        
+        current_section = None
+        for line in context.split('\n'):
+            line = line.strip()
+            if ':' in line:
+                key, value = line.split(':', 1)
+                key = key.strip().lower()
+                if key == 'immediate actions':
+                    current_section = 'advice'
+                    components['advice'].append(value.strip())
+                elif key == 'red flags':
+                    current_section = 'red_flags'
+                    components['red_flags'].append(value.strip())
+                elif key == 'action timeline':
+                    components['timeframe'] = value.strip()
+                elif key == 'cultural considerations':
+                    current_section = 'special_considerations'
+                    components['special_considerations'].append(value.strip())
+            elif current_section:
+                components[current_section].append(line)
+        
+        # Convert lists to strings
+        for key in components:
+            if isinstance(components[key], list):
+                components[key] = '\n- '.join(components[key])
+        return components
+
     def setup_model(self, model_name="google/gemma-7b-it"):
         """Initialize the medical AI model with 4-bit quantization"""
         if self.model is not None:
@@ -121,34 +178,211 @@ class MedicalTriageBot:
 
     def build_medical_index(self):
         medical_knowledge = {
-            "emergency_protocols.txt": """Emergency Protocols:
-            - Chest pain: Call 999 immediately
-            - Breathing difficulties: Urgent 999 call
-            - Severe bleeding: Apply pressure, call 999""",
-            
-            "triage_guidelines.txt": """Triage Guidelines:
-            - Persistent fever >48h: Contact 111
-            - Minor injuries: Visit urgent care
-            - Medication questions: Consult GP""",
-            "gp_care.txt": """GP Management:
-            - Persistent cough >3 weeks: Book GP
-            - Medication reviews: Schedule appointment
-            - Chronic condition management""",
-            
-            "urgent_care.txt": """Urgent Care:
-            - Minor burns: Visit urgent care
-            - Sprains: Same-day urgent care
-            - Ear infections: Walk-in centre"""
+            # ===== EMERGENCY CONDITIONS (Call 999) =====
+            "emergency_pediatric.txt": """
+            Triage Level: Emergency
+            Conditions: Pediatric Respiratory Distress, Head Injury
+            Key Symptoms:
+            - Blue lips/tongue (cyanosis)
+            - Stridor or wheezing at rest
+            - Repeated vomiting post-head injury
+            - Unconsciousness >1 minute
+            Immediate Actions:
+            1. Call 999 immediately
+            2. Maintain airway if trained
+            3. Monitor breathing rate
+            Cultural Considerations:
+            - Provide gender-matched responders if requested
+            - Accommodate religious head coverings during assessment""",
+    
+            "emergency_mental_health.txt": """
+            Triage Level: Emergency
+            Conditions: Suicidal Crisis, Psychotic Episode
+            Key Symptoms:
+            - Expressing concrete suicide plan
+            - Hallucinations commanding harm
+            - Severe self-injury with blood loss
+            - Catatonic state
+            Immediate Actions:
+            1. Call 999 if immediate danger
+            2. Remove potential weapons
+            3. Stay with patient until help arrives
+            Special Considerations:
+            - Avoid physical restraint unless absolutely necessary
+            - Use non-judgmental language
+            Red Flags:
+            - Talking about being a burden
+            - Sudden calm after depression
+            - Giving away possessions""",
+    
+            # ===== URGENT CARE CONDITIONS (A&E/UTC) =====
+            "urgent_mental_health.txt": """
+            Triage Level: Urgent
+            Conditions: Panic Attacks, Self-Harm
+            Key Symptoms:
+            - Hyperventilation >30 minutes
+            - Superficial self-harm injuries
+            - Acute anxiety with derealization
+            Action Timeline: Within 4 hours
+            While Waiting:
+            - Practice grounding techniques
+            - Use ice cubes for sensory focus
+            - Track panic duration/frequency
+            Escalate If:
+            - Chest pain develops
+            - Dissociation persists >1 hour
+            - Urges escalate""",
+    
+            "urgent_pregnancy.txt": """
+            Triage Level: Urgent
+            Conditions: Pregnancy Complications
+            Key Symptoms:
+            - Vaginal bleeding + abdominal pain
+            - Reduced fetal movement
+            - Sudden swelling + headache
+            Action Timeline: Within 2 hours
+            Cultural Protocols:
+            - Offer female clinician if preferred
+            - Respect modesty requests
+            Red Flags:
+            - Fluid leakage + fever
+            - Visual disturbances
+            - Contractions <37 weeks""",
+    
+            # ===== GP CARE CONDITIONS =====
+            "gp_mental_health.txt": """
+            Triage Level: GP Care
+            Conditions: Anxiety, Depression, PTSD
+            Key Symptoms:
+            - Persistent low mood >2 weeks
+            - Panic attacks 2+/week
+            - Sleep disturbances + fatigue
+            - Avoidance behaviors
+            Action Timeline: 72 hours
+            Management Strategies:
+            - Keep mood/symptom diary
+            - Practice 4-7-8 breathing
+            - Maintain routine activities
+            Red Flags:
+            - Social withdrawal >1 week
+            - Weight loss >5% in month
+            - Suicidal ideation""",
+    
+            # ===== SELF-CARE CONDITIONS =====
+            "selfcare_mental_health.txt": """
+            Triage Level: Self-Care
+            Conditions: Mild Anxiety, Stress
+            Key Symptoms:
+            - Situational anxiety
+            - Work-related stress
+            - Mild sleep difficulties
+            Management:
+            - Practice box breathing
+            - Limit caffeine/alcohol
+            - Use worry journal
+            - Progressive muscle relaxation
+            Escalate If:
+            - Symptoms persist >2 weeks
+            - Panic attacks develop
+            - Daily functioning impaired""",
+    
+            # ===== PEDIATRIC MENTAL HEALTH =====
+            "pediatric_mental_health.txt": """
+            Triage Level: GP Care
+            Conditions: Childhood Anxiety, School Refusal
+            Key Symptoms:
+            - School avoidance >3 days
+            - Somatic complaints (stomachaches)
+            - Nightmares/bedwetting regression
+            Action Timeline: 1 week
+            Parent Guidance:
+            - Maintain consistent routine
+            - Validate feelings without reinforcement
+            - Use gradual exposure techniques
+            Red Flags:
+            - Food restriction
+            - Self-harm marks
+            - Social isolation >1 week""",
+    
+            # ===== CULTURAL MENTAL HEALTH =====
+            "cultural_mental_health.txt": """
+            Triage Level: GP Care
+            Conditions: Culturally-Specific Presentations
+            Key Symptoms:
+            - Somatic complaints without medical cause
+            - Religious preoccupation/guilt
+            - Migration-related stress
+            Cultural Considerations:
+            - Use trained interpreters
+            - Consider spiritual assessments
+            - Respect family hierarchy
+            Management:
+            - Community support referrals
+            - Culturally-adapted CBT
+            - Family involvement""",
+    
+            # ===== CHRONIC CONDITIONS =====
+            "chronic_mental_health.txt": """
+            Triage Level: GP Care
+            Conditions: Bipolar, OCD, Eating Disorders
+            Key Symptoms:
+            - Manic episodes >4 days
+            - Compulsions >1hr/day
+            - BMI <18.5 with body dysmorphia
+            Monitoring:
+            - Mood tracking charts
+            - Compulsion frequency log
+            - Weekly weight checks
+            Crisis Signs:
+            - Rapid speech + sleeplessness
+            - Food restriction >24hrs
+            - Contamination fears preventing eating""",
+    
+            # ===== PERINATAL MENTAL HEALTH =====
+            "perinatal_mental_health.txt": """
+            Triage Level: Urgent
+            Conditions: Postpartum Depression, Psychosis
+            Key Symptoms:
+            - Intrusive harm thoughts
+            - Detachment from baby
+            - Visual/auditory hallucinations
+            Action Timeline: 24 hours
+            Safety Measures:
+            - Partner supervision
+            - Remove harmful objects
+            - Breastfeeding support
+            Red Flags:
+            - Planning infant harm
+            - Extreme paranoia
+            - Refusing sleep""",
+    
+            # ===== ADDICTION & SUBSTANCE ===== 
+            "addiction_care.txt": """
+            Triage Level: Urgent
+            Conditions: Withdrawal, Overdose Risk
+            Key Symptoms:
+            - Seizure history + alcohol use
+            - IV drug use + fever
+            - Opioid constricted pupils
+            Immediate Actions:
+            1. Call 111 for withdrawal management
+            2. Monitor breathing rate
+            3. Provide naloxone if available
+            Danger Signs:
+            - Jaundice + abdominal pain
+            - Chest pain + stimulant use
+            - Hallucinations + tremor"""
         }
-        
+    
+        # Keep existing index building logic
         os.makedirs("medical_knowledge", exist_ok=True)
         for filename, content in medical_knowledge.items():
             with open(f"medical_knowledge/{filename}", "w") as f:
                 f.write(content)
-        
+    
         text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=self.chunk_size,
-            chunk_overlap=self.chunk_overlap
+            chunk_size=512,
+            chunk_overlap=128
         )
         
         documents = []
@@ -173,68 +407,51 @@ class MedicalTriageBot:
     @torch.inference_mode()
     def generate_safe_response(self, message, history):
         try:
+            # Add case initialization check
             if self.current_case is None:
                 context = self.get_medical_context(message)
                 self._initialize_case(context)
                 return self._next_question()
+                
+            context = self.get_medical_context(message)
+            triage_level = self._determine_triage_level(context)
+            components = self._parse_medical_context(context)
             
-            # Existing state handling code
-            if self.current_case["stage"] == "investigation":
-                return self._handle_investigation(message, history)
-            
-            # Add model inference HERE
-            inputs = self.tokenizer(
-                self._build_prompt(message, history),
-                return_tensors="pt",
-                truncation=True,
-                max_length=1024
-            ).to(self.model.device)
-            
-            # Create safety-focused prompt
-            prompt = f"""<start_of_turn>system
-You are a medical triage assistant. Use this context:
-{context}
-Current conversation:
-{conversation}
-Guidelines:
-1. Assess symptom severity
-2. Recommend appropriate care level
-3. Never diagnose or prescribe
-4. Always include safety netting
-<end_of_turn>
-<start_of_turn>user
-{message}
-<end_of_turn>
-<start_of_turn>assistant"""
-            
-            # Generate response
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                truncation=True,
-                max_length=1024
-            ).to(self.model.device)
-            
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=256,
-                temperature=0.7,
-                top_p=0.9,
-                repetition_penalty=1.2
-            )
+            response_templates = {
+                'Emergency': (
+                    "🚨 **Emergency Alert**\n"
+                    "{advice}\n\n"
+                    "⚠️ **Immediate Action Required:**\n"
+                    "- Call 999 NOW if:\n{red_flags}\n"
+                    "🌍 **Special Considerations:**\n{special_considerations}"
+                ),
+                'Urgent': (
+                    "⚠️ **Urgent Care Needed**\n"
+                    "Visit A&E within {timeframe} if:\n{red_flags}\n\n"
+                    "🩺 **While Waiting:**\n{advice}\n\n"
+                    "🌍 **Cultural Notes:**\n{special_considerations}"
+                ),
+                'GP Care': (
+                    "📅 **GP Consultation**\n"
+                    "Book appointment within {timeframe}\n\n"
+                    "💡 **Self-Care Advice:**\n{advice}\n\n"
+                    "⚠️ **Red Flags:**\n{red_flags}"
+                ),
+                'Self-Care': (
+                    "🏡 **Self-Care Management**\n"
+                    "{advice}\n\n"
+                    "⚠️ **Seek Help If:**\n{red_flags}"
+                )
+            }
             
-            # Clean and validate response
-            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            clean_response = full_response.split("<start_of_turn>assistant")[-1]
-            clean_response = clean_response.split("<end_of_turn>")[0].strip()
+            response = response_templates[triage_level].format(**components)
             
-            # Ensure medical safety
-            if not any(phrase in clean_response.lower() for phrase in self.safety_phrases):
-                clean_response += "\n\nIf symptoms persist, please contact NHS 111."
-            
-            self.last_interaction_time = time.time()
-            return clean_response[:500]  # Limit response length
+            # Add safety netting
+            if not any(phrase in response.lower() for phrase in self.safety_phrases):
+                response += "\n\nIf symptoms persist, please contact NHS 111."
             
+            return response[:500]  # Maintain length limit
+    
         except Exception as e:
             logger.error(f"Generation error: {e}")
             return "Please contact NHS 111 directly for urgent medical advice."
@@ -311,20 +528,25 @@ Guidelines:
     
     def _get_location(self, postcode):
         return "London" if postcode.startswith("L") else "Manchester" 
-
-def _build_prompt(self, message, history):
-    conversation = "\n".join([f"User: {user}\nAssistant: {bot}" for user, bot in history[-3:]])
-    context = self.get_medical_context(message)
+        
     
-    return f"""<start_of_turn>system
+
+    def _build_prompt(self, message, history):
+        conversation = "\n".join([f"User: {user}\nAssistant: {bot}" for user, bot in history[-3:]])
+        context = self.get_medical_context(message)
+        triage_level = self._determine_triage_level(context)
+        
+        return f"""<start_of_turn>system
+Triage Level: {triage_level}
 Context:
 {context}
-Conversation:
+Conversation History:
 {conversation}
-Guidelines:
-1. Follow investigation flow
-2. Consider location: {self.current_case.get('location', 'unknown')}
-3. Maintain safety protocols
+Response Guidelines:
+1. Use {triage_level} response template
+2. Include safety netting
+3. Consider cultural factors
+4. Maintain NHS protocols
 <end_of_turn>
 <start_of_turn>user
 {message}
@@ -340,10 +562,6 @@ def create_medical_interface():
     bot.setup_model()
     bot.setup_rag_system()
 
-    def create_medical_interface():
-        bot = MedicalTriageBot()
-        bot.setup_model()
-        bot.setup_rag_system()
 
     def handle_conversation(message, history):
         try: