Update app.py

app.py

@@ -2,96 +2,433 @@ 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
 
-MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
+# Enhanced State Management
+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  # intake, assessment, summary, recommendations
 
-SYSTEM_PROMPT = """You are a professional virtual doctor. Your goal is to collect detailed information about the user's health condition, symptoms, medical history, medications, lifestyle, and other relevant data.
+# Medical Knowledge Base
+MEDICAL_CATEGORIES = {
+    "respiratory": ["cough", "shortness of breath", "chest pain", "wheezing"],
+    "gastrointestinal": ["nausea", "vomiting", "diarrhea", "stomach pain", "heartburn"],
+    "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"]
+}
 
-Ask 1-2 follow-up questions at a time to gather more details about:
-- Detailed description of symptoms
-- Duration (when did it start?)
-- Severity (scale of 1-10)
-- Aggravating or alleviating factors
-- Related symptoms
-- Medical history
-- Current medications and allergies
+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"
+]
 
-After collecting sufficient information (4-5 exchanges), summarize findings and suggest when they should seek professional care. Do NOT make specific diagnoses or recommend specific treatments.
+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?"
+}
 
-Respond empathetically and clearly. Always be professional and thorough."""
-
-print("Loading model...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
-print("Model loaded successfully!")
-
-# Conversation state tracking
-conversation_turns = {}
-
-def build_llama2_prompt(system_prompt, history, user_input):
-    """Format the conversation history and user input for Llama-2 chat models."""
-    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
+class EnhancedMedicalAssistant:
+    def __init__(self):
+        self.load_models()
+        self.setup_langgraph()
+    
+    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!")
+    
+    def setup_langgraph(self):
+        """Setup 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"
+            }
+        )
+        workflow.add_edge("generate_recommendations", END)
+        workflow.add_edge("emergency_triage", END)
+        
+        self.workflow = workflow.compile()
     
-    # Add conversation history
-    for user_msg, assistant_msg in history:
-        prompt += f"{user_msg} [/INST] {assistant_msg} </s><s>[INST] "
+    def patient_intake(self, state: MedicalState) -> MedicalState:
+        """Initial patient intake and basic information gathering"""
+        last_message = state["conversation_history"][-1]["content"] if state["conversation_history"] else ""
+        
+        # Extract symptoms and categorize them
+        detected_symptoms = self.extract_symptoms(last_message)
+        state["symptoms"].update(detected_symptoms)
+        
+        # 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
     
-    # Add the current user input
-    prompt += f"{user_input} [/INST] "
+    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["assessment_complete"] = True
+        return state
     
-    return prompt
-
-@spaces.GPU
-def generate_response(message, history):
-    """Generate a response using the Llama-2 model with proper formatting."""
-    # Track conversation turns
-    session_id = "default-session"
-    if session_id not in conversation_turns:
-        conversation_turns[session_id] = 0
-    conversation_turns[session_id] += 1
-    
-    # Build the prompt with proper Llama-2 formatting
-    prompt = build_llama2_prompt(SYSTEM_PROMPT, history, message)
-    
-    # Add summarization instruction after 4 turns
-    if conversation_turns[session_id] >= 4:
-        prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
-    
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    # Generate the response
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs.input_ids,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+    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:
+            state["consultation_stage"] = "emergency"
+        else:
+            state["consultation_stage"] = "recommendations"
+        
+        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)
+        state["suggested_actions"] = recommendations
+        
+        return state
+    
+    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"
+        }
+        state["suggested_actions"] = [emergency_response]
+        return state
+    
+    def route_after_intake(self, state: MedicalState):
+        """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"
+    
+    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 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
+        state = MedicalState(
+            patient_id="session_001",
+            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"
         )
+        
+        # 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)
     
-    # Decode and extract only the assistant's response
-    full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
-    assistant_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
+    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.
+        """
     
-    return assistant_response
+    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}"
+        
+        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.
+        """
+        
+        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()
 
-# Create the Gradio interface
+# Initialize the enhanced medical assistant
+medical_assistant = EnhancedMedicalAssistant()
+
+@spaces.GPU
+def chat_interface(message, history):
+    """Gradio chat interface"""
+    return medical_assistant.generate_response(message, history)
+
+# Create Gradio interface
 demo = gr.ChatInterface(
-    fn=generate_response,
-    title="Medical Assistant Chatbot",
-    description="Ask about your symptoms and I'll help gather relevant information.",
+    fn=chat_interface,
+    title="🏥 Advanced Medical AI Assistant",
+    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.
+    
+    ⚠️ **Important**: I'm not a replacement for professional medical care. Always consult healthcare providers for serious concerns.
+    """,
     examples=[
-        "I have a cough and my throat hurts",
-        "I've been having headaches for a week",
-        "My stomach has been hurting since yesterday"
+        "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"
     ],
-    theme="soft"
+    theme="soft",
+    css="""
+    .message.user { background-color: #e3f2fd; }
+    .message.bot { background-color: #f1f8e9; }
+    """
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)