Update app.py

app.py

@@ -57,8 +57,20 @@ class MedicalTriageBot:
             "999", "111", "emergency", "GP", "NHS",
             "consult a doctor", "seek medical attention"
         ]
+        self.current_case = None
+        self.location_services = {
+            "London": {"A&E": ["St Thomas' Hospital", "Royal London Hospital"],
+                      "Urgent Care": ["UCLH Urgent Care Centre"]},
+            "Manchester": {"A&E": ["Manchester Royal Infirmary"]}
+        }
+        self.base_questions = [
+            ("duration", "How long have you experienced this?"),
+            ("severity", "On a scale of 1-10, how severe is it?"),
+            ("emergency_signs", "Any difficulty breathing, chest pain, or confusion?")
+        ]
 
     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:
             return
         
@@ -117,7 +129,16 @@ class MedicalTriageBot:
             "triage_guidelines.txt": """Triage Guidelines:
             - Persistent fever >48h: Contact 111
             - Minor injuries: Visit urgent care
-            - Medication questions: Consult GP"""
+            - 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"""
         }
         
         os.makedirs("medical_knowledge", exist_ok=True)
@@ -152,17 +173,22 @@ class MedicalTriageBot:
     @torch.inference_mode()
     def generate_safe_response(self, message, history):
         try:
-            # Rate limiting
-            if time.time() - self.last_interaction_time < self.interaction_cooldown:
-                time.sleep(self.interaction_cooldown)
+            if self.current_case is None:
+                context = self.get_medical_context(message)
+                self._initialize_case(context)
+                return self._next_question()
             
-            # Convert Gradio history to conversational format
-            conversation = "\n".join(
-                [f"User: {user}\nAssistant: {bot}" for user, bot in history[-3:]]
-            )
+            # Existing state handling code
+            if self.current_case["stage"] == "investigation":
+                return self._handle_investigation(message, history)
             
-            # Get medical context
-            context = self.get_medical_context(message)
+            # 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
@@ -213,6 +239,99 @@ Guidelines:
             logger.error(f"Generation error: {e}")
             return "Please contact NHS 111 directly for urgent medical advice."
 
+    def _initialize_case(self, message):
+        context = self.get_medical_context(message)
+        self.current_case = {
+            "symptoms": self._detect_symptoms(context),
+            "responses": {},
+            "stage": "investigation",
+            "step": 0,
+            "location": None,
+            "guidelines": context
+        }
+
+    def _detect_symptoms(self, context):
+        return list(set(
+            line.split(":")[0].strip() 
+            for line in context.split("\n") 
+            if ":" in line
+        ))
+
+    def _next_question(self):
+        questions = self.base_questions + self._generate_custom_questions()
+        if self.current_case["step"] < len(questions):
+            return questions[self.current_case["step"]][1]
+        self.current_case["stage"] = "location"
+        return "Could you share your UK postcode for local service recommendations?"
+
+    def _generate_custom_questions(self):
+        return [(f"custom_{i}", q.split(":")[1].strip()) 
+                for q in self.current_case["guidelines"].split("\n")
+                if "?" in q]
+
+    def _handle_investigation(self, message, history):
+        self.current_case["responses"][self.base_questions[self.current_case["step"]][0]] = message
+        self.current_case["step"] += 1
+        return self._next_question()
+
+    def _handle_location(self, message):
+        self.current_case["location"] = self._get_location(message)
+        self.current_case["stage"] = "recommendation"
+        return self._final_recommendation()
+
+    def _final_recommendation(self):
+        action = self._determine_action()
+        location_info = self._get_location_services()
+        self.current_case = None
+        return f"{action}\n\n{location_info}"
+
+    def _determine_action(self):
+        if self._is_emergency():
+            return "🆘 Call 999 immediately. I can stay on the line with you."
+        if self._needs_gp():
+            return "📅 Please book a GP appointment. Would you like me to help with that?"
+        return "🏥 Visit your nearest urgent care centre:"
+
+    def _get_location_services(self):
+        if not self.current_case["location"]:
+            return "Find local services: https://www.nhs.uk/service-search"
+        return "\n".join([
+            f"{service_type}: {', '.join(services)}" 
+            for service_type, services in 
+            self.location_services.get(self.current_case["location"], {}).items()
+        ])
+
+    def _is_emergency(self):
+        return any(keyword in self.current_case["guidelines"] 
+                for keyword in ["999", "emergency", "stroke"])
+    
+    def _needs_gp(self):
+        return any(keyword in self.current_case["guidelines"]
+                for keyword in ["GP", "appointment", "persistent"])
+    
+    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
+Context:
+{context}
+Conversation:
+{conversation}
+Guidelines:
+1. Follow investigation flow
+2. Consider location: {self.current_case.get('location', 'unknown')}
+3. Maintain safety protocols
+<end_of_turn>
+<start_of_turn>user
+{message}
+<end_of_turn>
+<start_of_turn>assistant"""
+       
+
 # ===========================
 # 💬 SAFE GRADIO INTERFACE
 # ===========================
@@ -221,10 +340,25 @@ 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:
+            # Handle GP booking requests
+            if "book gp" in message.lower():
+                return history + [(message, "Redirecting to GP booking system...")]
+            
+            # Handle location input
+            if any(word in message.lower() for word in ["postcode", "zip code", "location"]):
+                return history + [(message, "Please enter your UK postcode:")]
+            
+            # Normal symptom processing
             response = bot.generate_safe_response(message, history)
             return history + [(message, response)]
+        
         except Exception as e:
             logger.error(f"Conversation error: {e}")
             return history + [(message, "System error - please refresh the page")]
@@ -235,7 +369,7 @@ def create_medical_interface():
         
         with gr.Row():
             chatbot = gr.Chatbot(
-                value=[("", "Hello! I'm your NHS digital assistant. How can I help you today?")],
+                value=[("", "Hello! I'm Pearly, your digital assistant. How can I help you today?")],
                 height=500,
                 label="Medical Triage Chat"
             )
@@ -264,7 +398,7 @@ def create_medical_interface():
         ).then(lambda: "", None, [message_input])
         
         clear_btn.click(
-            lambda: [("", "Hello! I'm your NHS digital assistant. How can I help you today?")],
+            lambda: [("", "Hello! I'm Pearly, your digital assistant. How can I help you today?")],
             None,
             [chatbot]
         )