app.py

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
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 with RAG
class MedicalState(TypedDict):
    patient_id: str
    conversation_history: List[Dict]
    symptoms: Dict[str, any]
    vital_questions_asked: List[str]
    medical_history: Dict
    current_medications: List[str]
    allergies: List[str]
    severity_scores: Dict[str, int]
    red_flags: List[str]
    assessment_complete: bool
    suggested_actions: List[str]
    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_CATEGORIES = {
    "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"],
    "dermatological": ["rash", "itching", "skin changes", "wounds"],
    "mental_health": ["anxiety", "depression", "sleep issues", "stress"]
}

RED_FLAGS = [
    "chest pain", "difficulty breathing", "severe headache", "high fever",
    "blood in stool", "blood in urine", "severe abdominal pain",
    "sudden vision changes", "loss of consciousness", "severe allergic reaction"
]

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...")
        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 simplified LangGraph workflow"""
        workflow = StateGraph(MedicalState)
        
        workflow.add_node("intake", self.patient_intake)
        workflow.add_node("generate_recommendations", self.generate_recommendations)
        workflow.add_node("emergency_triage", self.emergency_triage)
        
        workflow.set_entry_point("intake")
        workflow.add_conditional_edges(
            "intake",
            self.route_after_intake,
            {
                "emergency": "emergency_triage",
                "recommendations": "generate_recommendations"
            }
        )
        workflow.add_edge("generate_recommendations", END)
        workflow.add_edge("emergency_triage", END)
        
        self.workflow = workflow.compile()
    
    def patient_intake(self, state: MedicalState) -> MedicalState:
        """Enhanced patient intake with RAG"""
        last_message = state["conversation_history"][-1]["content"] if state["conversation_history"] else ""
        
        # 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)
        state["red_flags"].extend(red_flags)
        
        # Determine consultation stage
        if red_flags:
            state["consultation_stage"] = "emergency"
        else:
            state["consultation_stage"] = "recommendations"
        
        return state
    
    def generate_recommendations(self, state: MedicalState) -> MedicalState:
        """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 = 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
    
    def route_after_intake(self, state: MedicalState):
        """Route decision after intake"""
        if state["red_flags"]:
            return "emergency"
        else:
            return "recommendations"
    
    def extract_symptoms(self, text: str) -> Dict:
        """Extract and categorize symptoms from patient text"""
        symptoms = {}
        text_lower = text.lower()
        
        for category, symptom_list in MEDICAL_CATEGORIES.items():
            for symptom in symptom_list:
                if symptom in text_lower:
                    symptoms[symptom] = {
                        "category": category,
                        "mentioned_at": datetime.now().isoformat(),
                        "context": text
                    }
        
        return symptoms
    
    def check_red_flags(self, text: str) -> List[str]:
        """Check for emergency red flags"""
        found_flags = []
        text_lower = text.lower()
        
        for flag in RED_FLAGS:
            if flag in text_lower:
                found_flags.append(flag)
        
        return found_flags
    
    def generate_response(self, message: str, history: List) -> str:
        """Main response generation function"""
        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_id,
            conversation_history=history + [{"role": "user", "content": message}],
            symptoms={},
            vital_questions_asked=[],
            medical_history={},
            current_medications=[],
            allergies=[],
            severity_scores={},
            red_flags=[],
            assessment_complete=False,
            suggested_actions=[],
            consultation_stage="intake",
            retrieved_knowledge=[],
            confidence_scores={}
        )
        
        # 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 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)
        
        # 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 generate_llama_response(self, prompt: str) -> str:
        """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 medical assistant
medical_assistant = EnhancedMedicalAssistant()

@spaces.GPU  
def chat_interface(message, history):
    """Gradio chat interface"""
    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="🏥 Medical AI Assistant with medRAG",
    description="""
    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 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; 
        border-radius: 10px;
        padding: 10px;
        margin: 5px;
    }
    .message.bot { 
        background-color: #f1f8e9; 
        border-radius: 10px;
        padding: 10px;
        margin: 5px;
    }
    """
)

if __name__ == "__main__":
    demo.launch(share=True)