Update app.py

app.py

@@ -6,8 +6,14 @@ from langgraph.graph import StateGraph, END
 from typing import TypedDict, List, Dict, Optional
 from datetime import datetime
 import json
+import re
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import faiss
+import pickle
+import os
 
-# Enhanced State Management
+# Enhanced State Management with RAG
 class MedicalState(TypedDict):
     patient_id: str
     conversation_history: List[Dict]
@@ -20,12 +26,49 @@ class MedicalState(TypedDict):
     red_flags: List[str]
     assessment_complete: bool
     suggested_actions: List[str]
-    consultation_stage: str  # intake, assessment, summary, recommendations
+    consultation_stage: str
+    retrieved_knowledge: List[Dict]
+    confidence_scores: Dict[str, float]
+
+# Medical Knowledge Base for RAG
+MEDICAL_KNOWLEDGE_BASE = {
+    "conditions": {
+        "common_cold": {
+            "symptoms": ["runny nose", "cough", "sneezing", "sore throat", "mild fever"],
+            "treatment": "Rest, fluids, OTC pain relievers",
+            "otc_medications": [
+                {"name": "Acetaminophen", "dose": "500-1000mg every 4-6 hours", "max_daily": "3000mg"},
+                {"name": "Ibuprofen", "dose": "200-400mg every 4-6 hours", "max_daily": "1200mg"}
+            ],
+            "home_remedies": ["Warm salt water gargle", "Honey and lemon tea", "Steam inhalation"],
+            "when_to_seek_care": "If symptoms worsen after 7-10 days or fever above 101.3°F"
+        },
+        "headache": {
+            "symptoms": ["head pain", "pressure", "throbbing"],
+            "treatment": "Pain relief, rest, hydration",
+            "otc_medications": [
+                {"name": "Acetaminophen", "dose": "500-1000mg every 4-6 hours", "max_daily": "3000mg"},
+                {"name": "Ibuprofen", "dose": "400-600mg every 6-8 hours", "max_daily": "1200mg"}
+            ],
+            "home_remedies": ["Cold or warm compress", "Rest in dark room", "Stay hydrated"],
+            "when_to_seek_care": "Sudden severe headache, fever, neck stiffness, vision changes"
+        },
+        "stomach_pain": {
+            "symptoms": ["abdominal pain", "nausea", "bloating", "cramps"],
+            "treatment": "Bland diet, rest, hydration",
+            "otc_medications": [
+                {"name": "Pepto-Bismol", "dose": "525mg every 30 minutes as needed", "max_daily": "8 doses"},
+                {"name": "TUMS", "dose": "2-4 tablets as needed", "max_daily": "15 tablets"}
+            ],
+            "home_remedies": ["BRAT diet", "Ginger tea", "Warm compress on stomach"],
+            "when_to_seek_care": "Severe pain, fever, vomiting, blood in stool"
+        }
+    }
+}
 
-# Medical Knowledge Base
 MEDICAL_CATEGORIES = {
-    "respiratory": ["cough", "shortness of breath", "chest pain", "wheezing"],
-    "gastrointestinal": ["nausea", "vomiting", "diarrhea", "stomach pain", "heartburn"],
+    "respiratory": ["cough", "shortness of breath", "chest pain", "wheezing", "runny nose", "sore throat"],
+    "gastrointestinal": ["nausea", "vomiting", "diarrhea", "stomach pain", "heartburn", "bloating"],
     "neurological": ["headache", "dizziness", "numbness", "tingling"],
     "musculoskeletal": ["joint pain", "muscle pain", "back pain", "stiffness"],
     "cardiovascular": ["chest pain", "palpitations", "swelling", "fatigue"],
@@ -39,71 +82,99 @@ RED_FLAGS = [
     "sudden vision changes", "loss of consciousness", "severe allergic reaction"
 ]
 
-VITAL_QUESTIONS = {
-    "symptom_onset": "When did your symptoms first start?",
-    "severity": "On a scale of 1-10, how severe would you rate your symptoms?",
-    "triggers": "What makes your symptoms better or worse?",
-    "associated_symptoms": "Are you experiencing any other symptoms?",
-    "medical_history": "Do you have any chronic medical conditions?",
-    "medications": "Are you currently taking any medications?",
-    "allergies": "Do you have any known allergies?"
-}
+class SimpleRAGSystem:
+    def __init__(self):
+        self.knowledge_base = MEDICAL_KNOWLEDGE_BASE
+        self.setup_simple_retrieval()
+    
+    def setup_simple_retrieval(self):
+        """Setup simple keyword-based retrieval system"""
+        self.symptom_to_condition = {}
+        
+        for condition, data in self.knowledge_base["conditions"].items():
+            for symptom in data["symptoms"]:
+                if symptom not in self.symptom_to_condition:
+                    self.symptom_to_condition[symptom] = []
+                self.symptom_to_condition[symptom].append(condition)
+    
+    def retrieve_relevant_conditions(self, symptoms: List[str]) -> List[Dict]:
+        """Retrieve relevant medical conditions based on symptoms"""
+        relevant_conditions = {}
+        
+        for symptom in symptoms:
+            symptom_lower = symptom.lower()
+            
+            # Direct match
+            if symptom_lower in self.symptom_to_condition:
+                for condition in self.symptom_to_condition[symptom_lower]:
+                    if condition not in relevant_conditions:
+                        relevant_conditions[condition] = self.knowledge_base["conditions"][condition]
+            
+            # Partial match
+            for kb_symptom, conditions in self.symptom_to_condition.items():
+                if symptom_lower in kb_symptom or kb_symptom in symptom_lower:
+                    for condition in conditions:
+                        if condition not in relevant_conditions:
+                            relevant_conditions[condition] = self.knowledge_base["conditions"][condition]
+        
+        return [{"condition": k, "data": v} for k, v in relevant_conditions.items()]
 
 class EnhancedMedicalAssistant:
     def __init__(self):
         self.load_models()
+        self.rag_system = SimpleRAGSystem()
         self.setup_langgraph()
+        self.conversation_count = {}
     
     def load_models(self):
         """Load the AI models"""
         print("Loading models...")
-        # Llama-2 for conversation
-        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
-        self.model = AutoModelForCausalLM.from_pretrained(
-            "meta-llama/Llama-2-7b-chat-hf",
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-        
-        # Meditron for medical suggestions
-        self.meditron_tokenizer = AutoTokenizer.from_pretrained("epfl-llm/meditron-7b")
-        self.meditron_model = AutoModelForCausalLM.from_pretrained(
-            "epfl-llm/meditron-7b",
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-        print("Models loaded successfully!")
+        try:
+            # Llama-2 for conversation
+            self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                "meta-llama/Llama-2-7b-chat-hf",
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            
+            # Meditron for medical suggestions
+            self.meditron_tokenizer = AutoTokenizer.from_pretrained("epfl-llm/meditron-7b")
+            if self.meditron_tokenizer.pad_token is None:
+                self.meditron_tokenizer.pad_token = self.meditron_tokenizer.eos_token
+                
+            self.meditron_model = AutoModelForCausalLM.from_pretrained(
+                "epfl-llm/meditron-7b",
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            print("Models loaded successfully!")
+        except Exception as e:
+            print(f"Error loading models: {e}")
+            # Fallback - use only one model
+            self.tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+            self.model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+            self.meditron_tokenizer = self.tokenizer
+            self.meditron_model = self.model
     
     def setup_langgraph(self):
-        """Setup LangGraph workflow"""
+        """Setup simplified LangGraph workflow"""
         workflow = StateGraph(MedicalState)
         
-        # Add nodes
         workflow.add_node("intake", self.patient_intake)
-        workflow.add_node("symptom_assessment", self.assess_symptoms)
-        workflow.add_node("risk_evaluation", self.evaluate_risks)
         workflow.add_node("generate_recommendations", self.generate_recommendations)
         workflow.add_node("emergency_triage", self.emergency_triage)
         
-        # Define edges
         workflow.set_entry_point("intake")
         workflow.add_conditional_edges(
             "intake",
             self.route_after_intake,
             {
-                "continue_assessment": "symptom_assessment",
                 "emergency": "emergency_triage",
-                "complete": "generate_recommendations"
-            }
-        )
-        workflow.add_edge("symptom_assessment", "risk_evaluation")
-        workflow.add_conditional_edges(
-            "risk_evaluation",
-            self.route_after_risk_eval,
-            {
-                "emergency": "emergency_triage",
-                "continue": "generate_recommendations",
-                "need_more_info": "symptom_assessment"
+                "recommendations": "generate_recommendations"
             }
         )
         workflow.add_edge("generate_recommendations", END)
@@ -112,59 +183,25 @@ class EnhancedMedicalAssistant:
         self.workflow = workflow.compile()
     
     def patient_intake(self, state: MedicalState) -> MedicalState:
-        """Initial patient intake and basic information gathering"""
+        """Enhanced patient intake with RAG"""
         last_message = state["conversation_history"][-1]["content"] if state["conversation_history"] else ""
         
-        # Extract symptoms and categorize them
+        # Extract symptoms
         detected_symptoms = self.extract_symptoms(last_message)
         state["symptoms"].update(detected_symptoms)
         
+        # Use RAG to get relevant medical knowledge
+        if detected_symptoms:
+            symptom_names = list(detected_symptoms.keys())
+            relevant_conditions = self.rag_system.retrieve_relevant_conditions(symptom_names)
+            state["retrieved_knowledge"] = relevant_conditions
+        
         # Check for red flags
         red_flags = self.check_red_flags(last_message)
-        if red_flags:
-            state["red_flags"].extend(red_flags)
-        
-        # Determine what vital questions still need to be asked
-        missing_questions = self.get_missing_vital_questions(state)
-        
-        if missing_questions and len(state["conversation_history"]) < 6:
-            state["consultation_stage"] = "intake"
-            return state
-        else:
-            state["consultation_stage"] = "assessment"
-            return state
-    
-    def assess_symptoms(self, state: MedicalState) -> MedicalState:
-        """Detailed symptom assessment"""
-        # Analyze symptom patterns and severity
-        for symptom, details in state["symptoms"].items():
-            if "severity" not in details:
-                # Need to ask about severity
-                state["consultation_stage"] = "assessment"
-                return state
+        state["red_flags"].extend(red_flags)
         
-        state["assessment_complete"] = True
-        return state
-    
-    def evaluate_risks(self, state: MedicalState) -> MedicalState:
-        """Evaluate patient risks and urgency"""
-        risk_score = 0
-        
-        # Check red flags
-        if state["red_flags"]:
-            risk_score += len(state["red_flags"]) * 3
-        
-        # Check severity scores
-        for severity in state["severity_scores"].values():
-            if severity >= 8:
-                risk_score += 2
-            elif severity >= 6:
-                risk_score += 1
-        
-        # Check symptom duration and progression
-        # (Implementation would analyze timeline)
-        
-        if risk_score >= 5:
+        # Determine consultation stage
+        if red_flags:
             state["consultation_stage"] = "emergency"
         else:
             state["consultation_stage"] = "recommendations"
@@ -172,23 +209,82 @@ class EnhancedMedicalAssistant:
         return state
     
     def generate_recommendations(self, state: MedicalState) -> MedicalState:
-        """Generate treatment recommendations and care suggestions"""
-        patient_summary = self.create_patient_summary(state)
-        
-        # Use Meditron for medical recommendations
-        recommendations = self.get_meditron_recommendations(patient_summary)
+        """Generate RAG-enhanced recommendations"""
+        # Create structured recommendations from RAG knowledge
+        recommendations = self.create_structured_recommendations(state)
         state["suggested_actions"] = recommendations
-        
         return state
     
+    def create_structured_recommendations(self, state: MedicalState) -> List[str]:
+        """Create structured recommendations using RAG knowledge"""
+        recommendations = []
+        
+        if not state["retrieved_knowledge"]:
+            recommendations.append("I need more specific information about your symptoms to provide targeted recommendations.")
+            return recommendations
+        
+        # Process each relevant condition
+        for knowledge_item in state["retrieved_knowledge"][:2]:  # Limit to top 2 conditions
+            condition = knowledge_item["condition"]
+            data = knowledge_item["data"]
+            
+            # Format condition information
+            condition_info = f"\n**Possible Condition: {condition.replace('_', ' ').title()}**\n"
+            
+            # Add medications
+            if "otc_medications" in data:
+                condition_info += "\n**💊 Over-the-Counter Medications:**\n"
+                for med in data["otc_medications"]:
+                    condition_info += f"• **{med['name']}**: {med['dose']} (Max daily: {med['max_daily']})\n"
+            
+            # Add home remedies
+            if "home_remedies" in data:
+                condition_info += "\n**🏠 Home Remedies:**\n"
+                for remedy in data["home_remedies"]:
+                    condition_info += f"• {remedy}\n"
+            
+            # Add when to seek care
+            if "when_to_seek_care" in data:
+                condition_info += f"\n**⚠️ Seek Medical Care If:** {data['when_to_seek_care']}\n"
+            
+            recommendations.append(condition_info)
+        
+        # Add general advice
+        recommendations.append("""
+**📋 General Recommendations:**
+• Monitor your symptoms for any changes
+• Stay hydrated and get adequate rest
+• Follow medication instructions carefully
+• Don't exceed recommended dosages
+
+**🚨 Emergency Warning Signs:**
+• Severe worsening of symptoms
+• High fever (>101.3°F/38.5°C)
+• Difficulty breathing
+• Severe pain
+• Signs of dehydration
+        """)
+        
+        return recommendations
+    
     def emergency_triage(self, state: MedicalState) -> MedicalState:
         """Handle emergency situations"""
-        emergency_response = {
-            "urgent_care_needed": True,
-            "recommended_action": "Seek immediate medical attention",
-            "reasons": state["red_flags"],
-            "instructions": "Go to the nearest emergency room or call emergency services"
-        }
+        emergency_response = f"""
+🚨 **URGENT MEDICAL ATTENTION NEEDED** 🚨
+
+Based on your symptoms, I strongly recommend seeking immediate medical care because you mentioned: {', '.join(state['red_flags'])}
+
+**Immediate Actions:**
+• Go to the nearest emergency room, OR
+• Call emergency services (911), OR  
+• Contact your doctor immediately
+
+**Why This is Urgent:**
+These symptoms can indicate serious conditions that require professional medical evaluation and treatment.
+
+⚠️ **Disclaimer:** This is not a medical diagnosis, but these symptoms warrant immediate professional assessment.
+        """
+        
         state["suggested_actions"] = [emergency_response]
         return state
     
@@ -196,19 +292,8 @@ class EnhancedMedicalAssistant:
         """Route decision after intake"""
         if state["red_flags"]:
             return "emergency"
-        elif len(state["vital_questions_asked"]) < 5:
-            return "continue_assessment"
         else:
-            return "complete"
-    
-    def route_after_risk_eval(self, state: MedicalState):
-        """Route decision after risk evaluation"""
-        if state["consultation_stage"] == "emergency":
-            return "emergency"
-        elif state["assessment_complete"]:
-            return "continue"
-        else:
-            return "need_more_info"
+            return "recommendations"
     
     def extract_symptoms(self, text: str) -> Dict:
         """Extract and categorize symptoms from patient text"""
@@ -237,62 +322,18 @@ class EnhancedMedicalAssistant:
         
         return found_flags
     
-    def get_missing_vital_questions(self, state: MedicalState) -> List[str]:
-        """Determine which vital questions haven't been asked"""
-        asked = state["vital_questions_asked"]
-        return [q for q in VITAL_QUESTIONS.keys() if q not in asked]
-    
-    def create_patient_summary(self, state: MedicalState) -> str:
-        """Create a comprehensive patient summary"""
-        summary = f"""
-        Patient Summary:
-        Symptoms: {json.dumps(state['symptoms'], indent=2)}
-        Medical History: {state['medical_history']}
-        Current Medications: {state['current_medications']}
-        Allergies: {state['allergies']}
-        Severity Scores: {state['severity_scores']}
-        Conversation History: {[msg['content'] for msg in state['conversation_history'][-3:]]}
-        """
-        return summary
-    
-    def get_meditron_recommendations(self, patient_summary: str) -> List[str]:
-        """Get medical recommendations using Meditron model"""
-        prompt = f"""
-        Based on the following patient information, provide:
-        1. Specific over-the-counter medications with dosing
-        2. Home remedies and self-care measures
-        3. When to seek professional medical care
-        4. Follow-up recommendations
-        
-        Patient Information:
-        {patient_summary}
-        
-        Response:"""
-        
-        inputs = self.meditron_tokenizer(prompt, return_tensors="pt").to(self.meditron_model.device)
-        
-        with torch.no_grad():
-            outputs = self.meditron_model.generate(
-                inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=400,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True
-            )
-        
-        recommendation = self.meditron_tokenizer.decode(
-            outputs[0][inputs.input_ids.shape[1]:], 
-            skip_special_tokens=True
-        )
-        
-        return [recommendation]
-    
     def generate_response(self, message: str, history: List) -> str:
         """Main response generation function"""
-        # Initialize or update state
+        session_id = "default_session"
+        
+        # Track conversation count
+        if session_id not in self.conversation_count:
+            self.conversation_count[session_id] = 0
+        self.conversation_count[session_id] += 1
+        
+        # Initialize state
         state = MedicalState(
-            patient_id="session_001",
+            patient_id=session_id,
             conversation_history=history + [{"role": "user", "content": message}],
             symptoms={},
             vital_questions_asked=[],
@@ -303,130 +344,172 @@ class EnhancedMedicalAssistant:
             red_flags=[],
             assessment_complete=False,
             suggested_actions=[],
-            consultation_stage="intake"
+            consultation_stage="intake",
+            retrieved_knowledge=[],
+            confidence_scores={}
         )
         
-        # Run through LangGraph workflow
-        result = self.workflow.invoke(state)
-        
-        # Generate contextual response
-        response = self.generate_contextual_response(result, message)
-        
-        return response
-    
-    def generate_contextual_response(self, state: MedicalState, user_message: str) -> str:
-        """Generate a contextual response based on the current state"""
-        if state["consultation_stage"] == "emergency":
-            return self.format_emergency_response(state)
-        elif state["consultation_stage"] == "intake":
-            return self.format_intake_response(state, user_message)
-        elif state["consultation_stage"] == "assessment":
-            return self.format_assessment_response(state)
-        elif state["consultation_stage"] == "recommendations":
-            return self.format_recommendations_response(state)
-        else:
-            return self.format_default_response(user_message)
-    
-    def format_emergency_response(self, state: MedicalState) -> str:
-        """Format emergency response"""
-        return f"""
-        🚨 URGENT MEDICAL ATTENTION NEEDED 🚨
-        
-        Based on your symptoms, I recommend seeking immediate medical care because:
-        {', '.join(state['red_flags'])}
-        
-        Please:
-        - Go to the nearest emergency room, OR
-        - Call emergency services (911), OR
-        - Contact your doctor immediately
-        
-        This is not a diagnosis, but these symptoms warrant immediate professional evaluation.
-        """
+        # For first few messages, do conversational intake
+        if self.conversation_count[session_id] <= 3:
+            return self.generate_conversational_response(message, history)
+        
+        # After gathering info, run workflow for recommendations
+        try:
+            result = self.workflow.invoke(state)
+            return self.format_final_response(result)
+        except Exception as e:
+            print(f"Workflow error: {e}")
+            return self.generate_conversational_response(message, history)
     
-    def format_intake_response(self, state: MedicalState, user_message: str) -> str:
-        """Format intake response with follow-up questions"""
-        # Use Llama-2 to generate empathetic response
-        prompt = f"""
-        You are a caring virtual doctor. The patient said: "{user_message}"
+    def generate_conversational_response(self, message: str, history: List) -> str:
+        """Generate conversational response for intake phase"""
+        # Extract symptoms for context
+        symptoms = self.extract_symptoms(message)
+        red_flags = self.check_red_flags(message)
         
-        Respond empathetically and ask 1-2 specific follow-up questions about:
-        - Symptom details (duration, severity, triggers)
-        - Associated symptoms
-        - Medical history if relevant
-        
-        Be professional, caring, and thorough.
-        """
+        # Handle emergencies immediately
+        if red_flags:
+            return f"""
+🚨 **URGENT MEDICAL ATTENTION NEEDED** 🚨
+
+I notice you mentioned: {', '.join(red_flags)}
+
+Please seek immediate medical care:
+• Go to the nearest emergency room
+• Call emergency services (911)  
+• Contact your doctor immediately
+
+These symptoms require professional medical evaluation right away.
+            """
+        
+        # Generate contextual questions based on symptoms
+        if symptoms:
+            symptom_names = list(symptoms.keys())
+            prompt = f"""
+You are a caring medical assistant. The patient mentioned these symptoms: {', '.join(symptom_names)}.
+
+Respond empathetically and ask 1-2 relevant follow-up questions about:
+- How long they've had these symptoms
+- Severity (mild, moderate, severe)
+- What makes it better or worse
+- Any other symptoms they're experiencing
+
+Be professional, caring, and concise. Don't provide treatment advice yet.
+            """
+        else:
+            prompt = f"""
+You are a caring medical assistant. The patient said: "{message}"
+
+Respond empathetically and ask relevant questions to understand their health concern better.
+Be professional and caring.
+            """
         
         return self.generate_llama_response(prompt)
     
-    def format_assessment_response(self, state: MedicalState) -> str:
-        """Format detailed assessment response"""
-        return "Let me gather a bit more information to better understand your condition..."
-    
-    def format_recommendations_response(self, state: MedicalState) -> str:
-        """Format final recommendations"""
-        recommendations = "\n".join(state["suggested_actions"])
-        return f"""
-        Based on our consultation, here's my assessment and recommendations:
-        
-        {recommendations}
-        
-        **Important Disclaimer:** I am an AI assistant, not a licensed medical professional. 
-        These suggestions are for informational purposes only. Please consult with a 
-        healthcare provider for proper diagnosis and treatment.
-        """
-    
-    def format_default_response(self, user_message: str) -> str:
-        """Format default response"""
-        return self.generate_llama_response(f"Respond professionally to: {user_message}")
-    
     def generate_llama_response(self, prompt: str) -> str:
-        """Generate response using Llama-2"""
-        formatted_prompt = f"<s>[INST] {prompt} [/INST] "
-        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
-        
-        with torch.no_grad():
-            outputs = self.model.generate(
-                inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=300,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-        
-        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-        return response.split('</s>')[0].strip()
+        """Generate response using Llama-2 with better formatting"""
+        try:
+            formatted_prompt = f"<s>[INST] {prompt} [/INST]"
+            inputs = self.tokenizer(formatted_prompt, return_tensors="pt", truncation=True, max_length=512)
+            
+            if torch.cuda.is_available():
+                inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+            
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=200,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.eos_token_id
+                )
+            
+            # Decode response
+            response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+            
+            # Clean up the response
+            response = response.split('</s>')[0].strip()
+            response = response.replace('<s>', '').replace('[INST]', '').replace('[/INST]', '').strip()
+            
+            # Remove any XML-like tags
+            response = re.sub(r'<[^>]+>', '', response)
+            
+            return response if response else "I understand your concern. Can you tell me more about what you're experiencing?"
+            
+        except Exception as e:
+            print(f"Error generating response: {e}")
+            return "I understand your concern. Can you tell me more about your symptoms?"
+    
+    def format_final_response(self, state: MedicalState) -> str:
+        """Format the final response with recommendations"""
+        if state["consultation_stage"] == "emergency":
+            return state["suggested_actions"][0] if state["suggested_actions"] else "Please seek immediate medical attention."
+        
+        # Format recommendations nicely
+        if state["suggested_actions"]:
+            response = "## 🏥 Medical Assessment & Recommendations\n\n"
+            response += "Based on our conversation, here's what I recommend:\n"
+            
+            for action in state["suggested_actions"]:
+                response += f"{action}\n"
+            
+            response += "\n---\n"
+            response += "**Important Disclaimer:** I'm an AI assistant providing general health information. "
+            response += "This is not a substitute for professional medical advice, diagnosis, or treatment. "
+            response += "Always consult with qualified healthcare providers for medical concerns."
+            
+            return response
+        else:
+            return "Please provide more details about your symptoms so I can offer better guidance."
 
-# Initialize the enhanced medical assistant
+# Initialize the medical assistant
 medical_assistant = EnhancedMedicalAssistant()
 
-@spaces.GPU
+@spaces.GPU  
 def chat_interface(message, history):
     """Gradio chat interface"""
-    return medical_assistant.generate_response(message, history)
+    try:
+        return medical_assistant.generate_response(message, history)
+    except Exception as e:
+        print(f"Chat interface error: {e}")
+        return f"I apologize, but I encountered an error. Please try rephrasing your question. Error: {str(e)}"
 
 # Create Gradio interface
 demo = gr.ChatInterface(
     fn=chat_interface,
-    title="🏥 Advanced Medical AI Assistant",
+    title="🏥 Medical AI Assistant with medRAG",
     description="""
-    I'm an AI medical assistant that can help assess your symptoms and provide guidance. 
-    I'll ask relevant questions to better understand your condition and provide appropriate recommendations.
+    I'm an AI medical assistant powered by medical knowledge retrieval (medRAG). 
+    I can help assess your symptoms and provide evidence-based recommendations.
+    
+    **How it works:**
+    1. Tell me about your symptoms
+    2. I'll ask follow-up questions  
+    3. I'll provide personalized recommendations based on medical knowledge
     
     ⚠️ **Important**: I'm not a replacement for professional medical care. Always consult healthcare providers for serious concerns.
     """,
     examples=[
-        "I've been having severe chest pain for the last hour",
-        "I have a persistent cough that's been going on for 2 weeks",
-        "I'm experiencing nausea and stomach pain after eating",
-        "I have a headache and feel dizzy"
+        "I have a bad cough and sore throat",
+        "I've been having headaches for the past few days", 
+        "My stomach has been hurting after meals",
+        "I have chest pain and trouble breathing"
     ],
     theme="soft",
     css="""
-    .message.user { background-color: #e3f2fd; }
-    .message.bot { background-color: #f1f8e9; }
+    .message.user { 
+        background-color: #e3f2fd; 
+        border-radius: 10px;
+        padding: 10px;
+        margin: 5px;
+    }
+    .message.bot { 
+        background-color: #f1f8e9; 
+        border-radius: 10px;
+        padding: 10px;
+        margin: 5px;
+    }
     """
 )