Update app.py

app.py

@@ -1,306 +1,257 @@
 import gradio as gr
-import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import re
+from typing import Annotated, List, Dict, Any
+from typing_extensions import TypedDict
+from langgraph.graph import StateGraph, START
+from langgraph.graph.message import add_messages
 
 # Model configuration
 LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
 MEDITRON_MODEL = "epfl-llm/meditron-7b"
 
-SYSTEM_PROMPT = """You are a professional virtual doctor conducting a medical consultation. Your role is to gather comprehensive information about the patient's condition through intelligent questioning.
+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.
+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
+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.
+Respond empathetically and clearly. Always be professional and thorough."""
 
-**CONSULTATION APPROACH:**
-- Ask thoughtful, relevant follow-up questions based on the patient's responses
-- Prioritize gathering information about: symptom details, duration, severity, triggers, related symptoms, medical history, medications, and lifestyle factors
-- Ask 1-2 specific questions at a time that build naturally on their previous answers
-- Be empathetic, professional, and thorough in your questioning
-- Adapt your questions based on the symptoms they describe
-
-**IMPORTANT GUIDELINES:**
-- Generate intelligent follow-up questions that are contextually relevant to their specific symptoms
-- Don't ask generic questions - tailor each question to their particular situation
-- If they mention pain, ask about location, type, and triggers
-- If they mention duration, ask about progression or changes
-- Build each question logically from their previous responses
-
-After 4-5 meaningful exchanges, provide assessment and recommendations.
-Do NOT make specific prescriptions for prescription-only drugs.
-Always maintain a professional, caring tone throughout the consultation."""
-
-MEDITRON_PROMPT = """You are a board-certified physician providing evidence-based medical assessment. 
-
-Based on the patient information provided, please:
-1. Analyze the symptoms systematically
-2. Provide a differential diagnosis with most likely conditions
-3. Recommend appropriate next steps (testing, monitoring, or treatment)
-4. Suggest appropriate medications or remedies with dosing if applicable
-5. Include red flags that would require urgent medical attention
-6. Base recommendations on clinical guidelines
-
-Patient Information: {patient_info}
-
-Please provide a structured medical assessment:"""
+MEDITRON_PROMPT = """<|im_start|>system
+You are a specialized medical assistant focusing ONLY on suggesting over-the-counter medicines and home remedies based on patient information.
+Based on the following patient information, provide ONLY:
+1. One specific over-the-counter medicine with proper adult dosing instructions
+2. One practical home remedy that might help
+3. Clear guidance on when to seek professional medical care
+Be concise, practical, and focus only on general symptom relief. Do not diagnose. Include a disclaimer that you are not a licensed medical professional.
+<|im_end|>
+<|im_start|>user
+Patient information: {patient_info}
+<|im_end|>
+<|im_start|>assistant
+"""
 
-# Load models
-print("Loading models...")
-try:
-    tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        LLAMA_MODEL,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-    print("Llama-2 model loaded successfully!")
+print("Loading Llama-2 model...")
+tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
     
-    meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
-    if meditron_tokenizer.pad_token is None:
-        meditron_tokenizer.pad_token = meditron_tokenizer.eos_token
-        
-    meditron_model = AutoModelForCausalLM.from_pretrained(
-        MEDITRON_MODEL,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-    print("Meditron model loaded successfully!")
-except Exception as e:
-    print(f"Error loading models: {e}")
+model = AutoModelForCausalLM.from_pretrained(
+    LLAMA_MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+print("Llama-2 model loaded successfully!")
 
-class MedicalConsultationBot:
-    def __init__(self):
-        self.reset_conversation()
+print("Loading Meditron model...")
+meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
+if meditron_tokenizer.pad_token is None:
+    meditron_tokenizer.pad_token = meditron_tokenizer.eos_token
     
-    def reset_conversation(self):
-        """Reset all conversation state"""
-        self.conversation_history = []
-        self.patient_name = None
-        self.patient_age = None
-        self.medical_turns = 0
-        self.stage = "greeting"  # greeting -> name -> age -> symptoms -> diagnosis
-        
-    def add_to_history(self, user_message, bot_response):
-        """Add exchange to conversation history"""
-        self.conversation_history.append({
-            "user": user_message,
-            "bot": bot_response
-        })
-    
-    def get_conversation_context(self):
-        """Get full conversation context as string"""
-        context = ""
-        if self.patient_name:
-            context += f"Patient Name: {self.patient_name}\n"
-        if self.patient_age:
-            context += f"Patient Age: {self.patient_age}\n"
-        
-        context += "\nConversation History:\n"
-        for exchange in self.conversation_history:
-            context += f"Patient: {exchange['user']}\n"
-            context += f"Doctor: {exchange['bot']}\n"
-        
-        return context
-    
-    def build_llama_prompt(self, current_message):
-        """Build prompt for Llama model with conversation context"""
-        prompt = f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n"
-        
-        # Add conversation context
-        context = self.get_conversation_context()
-        if context.strip():
-            prompt += f"Previous conversation context:\n{context}\n\n"
-        
-        prompt += f"Current patient message: {current_message}\n\nProvide a professional medical response with appropriate follow-up questions. [/INST]"
-        
-        return prompt
-
-# Global bot instance
-medical_bot = MedicalConsultationBot()
+meditron_model = AutoModelForCausalLM.from_pretrained(
+    MEDITRON_MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+print("Meditron model loaded successfully!")
 
-def get_meditron_diagnosis(patient_info):
-    """Use Meditron model to generate medical assessment"""
-    try:
-        prompt = MEDITRON_PROMPT.format(patient_info=patient_info)
-        inputs = meditron_tokenizer(
-            prompt, 
-            return_tensors="pt", 
-            max_length=512, 
-            truncation=True
-        ).to(meditron_model.device)
-        
-        with torch.no_grad():
-            outputs = meditron_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=meditron_tokenizer.pad_token_id
-            )
-        
-        response = meditron_tokenizer.decode(
-            outputs[0][inputs.input_ids.shape[1]:], 
-            skip_special_tokens=True
-        ).strip()
-        
-        return response
-    except Exception as e:
-        return f"Error generating medical assessment: {str(e)}"
+# Define the state for our LangGraph
+class ChatbotState(TypedDict):
+    messages: Annotated[List, add_messages]
+    turn_count: int
+    patient_info: List[str]
 
-@spaces.GPU
-def medical_chat_response(message, history):
-    """Main chat response function with proper state management"""
-    global medical_bot
+# Function to build Llama-2 prompt
+def build_llama2_prompt(messages):
+    """Format the conversation history for Llama-2 chat models."""
+    prompt = f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n"
     
-    # If this is a new conversation (empty history), reset the bot
-    if not history:
-        medical_bot.reset_conversation()
+    # Add conversation history
+    for i, msg in enumerate(messages[:-1]):
+        if i % 2 == 0:  # User message
+            prompt += f"{msg.content} [/INST] "
+        else:  # Assistant message
+            prompt += f"{msg.content} </s><s>[INST] "
+    
+    # Add the current user input
+    prompt += f"{messages[-1].content} [/INST] "
     
-    user_message = message.strip()
+    return prompt
+
+# Function to get Llama-2 response
+def get_llama2_response(prompt, turn_count):
+    """Generate response from Llama-2 model."""
+    # Add summarization instruction after 4 turns
+    if turn_count >= 4:
+        prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
     
-    # Stage 1: Initial greeting and ask for name
-    if medical_bot.stage == "greeting":
-        bot_response = "Hello! I'm your AI medical assistant. Before we discuss your health concerns, could you please tell me your name?"
-        medical_bot.stage = "name"
-        medical_bot.add_to_history(user_message, bot_response)
-        return bot_response
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    # Stage 2: Collect name and ask for age
-    elif medical_bot.stage == "name":
-        medical_bot.patient_name = user_message
-        bot_response = f"Nice to meet you, {medical_bot.patient_name}! Could you please tell me your age?"
-        medical_bot.stage = "age"
-        medical_bot.add_to_history(user_message, bot_response)
-        return bot_response
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id
+        )
     
-    # Stage 3: Collect age and start medical consultation
-    elif medical_bot.stage == "age":
-        medical_bot.patient_age = user_message
-        bot_response = f"Thank you, {medical_bot.patient_name}! Now, what brings you here today? Please tell me about any symptoms or health concerns you're experiencing."
-        medical_bot.stage = "symptoms"
-        medical_bot.add_to_history(user_message, bot_response)
-        return bot_response
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
     
-    # Stage 4: Medical consultation - gather symptoms with intelligent follow-ups
-    elif medical_bot.stage == "symptoms":
-        medical_bot.medical_turns += 1
-        
-        # If we've had enough turns, move to diagnosis
-        if medical_bot.medical_turns >= 4:
-            medical_bot.stage = "diagnosis"
-            return generate_final_diagnosis(user_message)
-        
-        # Generate intelligent follow-up questions
-        try:
-            prompt = medical_bot.build_llama_prompt(user_message)
-            inputs = tokenizer(
-                prompt, 
-                return_tensors="pt", 
-                max_length=1024, 
-                truncation=True
-            ).to(model.device)
-            
-            with torch.no_grad():
-                outputs = model.generate(
-                    inputs.input_ids,
-                    attention_mask=inputs.attention_mask,
-                    max_new_tokens=200,
-                    temperature=0.8,
-                    top_p=0.95,
-                    do_sample=True,
-                    pad_token_id=tokenizer.pad_token_id
-                )
-            
-            full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            bot_response = full_response.split('[/INST]')[-1].strip()
-            
-            # Clean up the response
-            bot_response = bot_response.replace('<s>', '').replace('</s>', '').strip()
-            
-            medical_bot.add_to_history(user_message, bot_response)
-            return bot_response
-            
-        except Exception as e:
-            bot_response = f"I understand. Could you tell me more about how long you've been experiencing this and if there are any specific triggers or patterns you've noticed?"
-            medical_bot.add_to_history(user_message, bot_response)
-            return bot_response
+    return response
+
+# Function to get Meditron suggestions
+def get_meditron_suggestions(patient_info):
+    """Generate medicine and remedy suggestions from Meditron model."""
+    prompt = MEDITRON_PROMPT.format(patient_info=patient_info)
+    inputs = meditron_tokenizer(prompt, return_tensors="pt").to(meditron_model.device)
     
-    # Stage 5: Final diagnosis and treatment recommendations
-    elif medical_bot.stage == "diagnosis":
-        return generate_final_diagnosis(user_message)
+    with torch.no_grad():
+        outputs = meditron_model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=256,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=meditron_tokenizer.pad_token_id
+        )
     
-    # Handle any questions after diagnosis
-    else:
-        # Check if they're asking about their name or previous information
-        if "name" in user_message.lower() and medical_bot.patient_name:
-            return f"Your name is {medical_bot.patient_name}."
-        elif "age" in user_message.lower() and medical_bot.patient_age:
-            return f"You told me you are {medical_bot.patient_age} years old."
-        else:
-            return "Is there anything else about your health concerns I can help you with today?"
+    suggestion = meditron_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    return suggestion
 
-def generate_final_diagnosis(current_message):
-    """Generate final diagnosis using both models"""
-    global medical_bot
+# Define LangGraph nodes
+def process_user_input(state: ChatbotState) -> ChatbotState:
+    """Process user input and update state."""
+    # Extract the latest user message
+    user_message = state["messages"][-1].content
     
-    # Add current message to history
-    medical_bot.add_to_history(current_message, "")
+    # Update patient info
+    return {
+        "patient_info": state["patient_info"] + [user_message],
+        "turn_count": state["turn_count"] + 1
+    }
+
+def generate_llama_response(state: ChatbotState) -> ChatbotState:
+    """Generate response using Llama-2 model."""
+    prompt = build_llama2_prompt(state["messages"])
+    response = get_llama2_response(prompt, state["turn_count"])
     
-    # Compile complete patient information
-    patient_info = f"""
-Patient Name: {medical_bot.patient_name}
-Patient Age: {medical_bot.patient_age}
+    return {"messages": [{"role": "assistant", "content": response}]}
 
-Complete Consultation History:
-"""
+def check_turn_count(state: ChatbotState) -> str:
+    """Check if we need to add medicine suggestions."""
+    if state["turn_count"] >= 4:
+        return "add_suggestions"
+    return "continue"
+
+def add_medicine_suggestions(state: ChatbotState) -> ChatbotState:
+    """Add medicine suggestions from Meditron model."""
+    # Get the last assistant response
+    last_response = state["messages"][-1].content
     
-    for exchange in medical_bot.conversation_history[:-1]:  # Exclude the empty last entry
-        patient_info += f"Doctor: {exchange['bot']}\n"
-        patient_info += f"Patient: {exchange['user']}\n"
+    # Collect full patient conversation
+    full_patient_info = "\n".join(state["patient_info"]) + "\n\nSummary: " + last_response
     
-    patient_info += f"Patient: {current_message}\n"
+    # Get medicine suggestions
+    medicine_suggestions = get_meditron_suggestions(full_patient_info)
     
-    # Get diagnosis from Meditron
-    meditron_assessment = get_meditron_diagnosis(patient_info)
+    # Format final response
+    final_response = (
+        f"{last_response}\n\n"
+        f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
+        f"{medicine_suggestions}"
+    )
     
-    # Generate comprehensive response
-    final_response = f"""Thank you for providing all this information, {medical_bot.patient_name}. Based on our consultation, here is my assessment:
+    # Return updated message
+    return {"messages": [{"role": "assistant", "content": final_response}]}
 
-**MEDICAL ASSESSMENT & RECOMMENDATIONS:**
-
-{meditron_assessment}
-
-**IMPORTANT DISCLAIMER:** This assessment is for informational purposes only and should not replace professional medical advice. Please consult with a healthcare provider for proper diagnosis and treatment.
+# Build the LangGraph
+def build_graph():
+    """Build and return the LangGraph for our chatbot."""
+    graph = StateGraph(ChatbotState)
+    
+    # Add nodes
+    graph.add_node("process_input", process_user_input)
+    graph.add_node("generate_response", generate_llama_response)
+    graph.add_node("add_suggestions", add_medicine_suggestions)
+    
+    # Add edges
+    graph.add_edge(START, "process_input")
+    graph.add_edge("process_input", "generate_response")
+    graph.add_conditional_edges(
+        "generate_response",
+        check_turn_count,
+        {
+            "add_suggestions": "add_suggestions",
+            "continue": END
+        }
+    )
+    graph.add_edge("add_suggestions", END)
+    
+    return graph.compile()
 
-**NEXT STEPS:** I recommend scheduling an appointment with your primary care physician or appropriate specialist for further evaluation and personalized treatment.
+# Initialize the graph
+chatbot_graph = build_graph()
 
-Is there anything specific about this assessment you'd like me to clarify?"""
+# Function for Gradio interface
+def chat_response(message, history):
+    """Generate chatbot response using LangGraph."""
+    # Initialize state if this is the first message
+    if not history:
+        state = {
+            "messages": [{"role": "user", "content": message}],
+            "turn_count": 0,
+            "patient_info": []
+        }
+    else:
+        # Convert history to messages format
+        messages = []
+        for user_msg, bot_msg in history:
+            messages.append({"role": "user", "content": user_msg})
+            messages.append({"role": "assistant", "content": bot_msg})
+        
+        # Add current message
+        messages.append({"role": "user", "content": message})
+        
+        # Get turn count from history
+        turn_count = len(history)
+        
+        # Build patient info from history
+        patient_info = [user_msg for user_msg, _ in history]
+        
+        state = {
+            "messages": messages,
+            "turn_count": turn_count,
+            "patient_info": patient_info
+        }
     
-    # Update conversation history with final response
-    medical_bot.conversation_history[-1]["bot"] = final_response
-    medical_bot.stage = "complete"
+    # Process through LangGraph
+    result = chatbot_graph.invoke(state)
     
-    return final_response
+    # Return the latest assistant message
+    return result["messages"][-1].content
 
-# Create Gradio interface
+# Create the Gradio interface
 demo = gr.ChatInterface(
-    fn=medical_chat_response,
-    title="🩺 AI Medical Assistant with Memory",
-    description="I'm an AI medical assistant that will remember our conversation. I'll first ask for your basic information, then gather details about your symptoms through intelligent follow-up questions, and finally provide a medical assessment.",
+    fn=chat_response,
+    title="Medical Assistant with LangGraph",
+    description="Tell me about your symptoms, and after gathering enough information, I'll suggest potential remedies.",
     examples=[
-        "Hello, I need medical help",
-        "I have a persistent cough",
-        "I've been having headaches",
-        "My stomach hurts"
+        "I have a cough and my throat hurts",
+        "I've been having headaches for a week",
+        "My stomach has been hurting since yesterday"
     ],
-    theme="soft",
-    retry_btn=None,
-    undo_btn=None,
-    clear_btn="🔄 Start New Consultation"
+    theme="soft"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()