Refactor app.py to streamline user information collection by removing redundant prompts for name and age. Implement a simple state tracking mechanism for improved conversation flow and enhance medical consultation process with structured follow-up questions.

app.py

@@ -9,9 +9,7 @@ import re
 LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
 MEDITRON_MODEL = "epfl-llm/meditron-7b"
 
-SYSTEM_PROMPT = """You are a professional virtual doctor. Your goal is to collect detailed information about the user's name, age, health condition, symptoms, medical history, medications, lifestyle, and other relevant data.
-
-Always begin by asking for the user's name and age if not already provided.
+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.
 
 **IMPORTANT** Ask 1-2 follow-up questions at a time to gather more details about:
 - Detailed description of symptoms
@@ -22,7 +20,7 @@ Always begin by asking for the user's name and age if not already provided.
 - Medical history
 - Current medications and allergies
 
-After collecting sufficient information (at least 4-5 exchanges, but continue up to 10 if the user keeps responding), summarize findings, provide a likely diagnosis (if possible), and suggest when they should seek professional care.
+After collecting sufficient information, summarize findings, provide a likely diagnosis (if possible), and suggest when they should seek professional care.
 
 If enough information is collected, provide a concise, general diagnosis and a practical over-the-counter medicine and home remedy suggestion.
 
@@ -67,33 +65,14 @@ meditron_model = AutoModelForCausalLM.from_pretrained(
 )
 print("Meditron model loaded successfully!")
 
-# Initialize LangChain memory
-memory = ConversationBufferMemory(return_messages=True)
-
-def build_llama2_prompt(system_prompt, messages, user_input, followup_stage=None):
-    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
-    for msg in messages:
-        if msg.type == "human":
-            prompt += f"{msg.content} [/INST] "
-        elif msg.type == "ai":
-            prompt += f"{msg.content} </s><s>[INST] "
-    
-    # Add a specific follow-up question if in followup stage
-    if followup_stage is not None:
-        followup_questions = [
-            "Can you describe your main symptoms in more detail? What exactly are you experiencing?",
-            "How long have you been experiencing these symptoms? When did they first start?",
-            "On a scale of 1-10, how would you rate the severity of your symptoms?",
-            "Have you noticed anything that makes your symptoms better or worse? Any triggers or relief factors?",
-            "Do you have any other related symptoms, such as fever, fatigue, nausea, or changes in appetite?"
-        ]
-        if followup_stage < len(followup_questions):
-            prompt += f"{followup_questions[followup_stage]} [/INST] "
-        else:
-            prompt += f"{user_input} [/INST] "
-    else:
-        prompt += f"{user_input} [/INST] "
-    return prompt
+# Simple conversation state tracking
+conversation_state = {
+    'name': None,
+    'age': None,
+    'medical_turns': 0,
+    'has_name': False,
+    'has_age': False
+}
 
 def get_meditron_suggestions(patient_info):
     """Use Meditron model to generate medicine and remedy suggestions."""
@@ -113,183 +92,141 @@ def get_meditron_suggestions(patient_info):
     suggestion = meditron_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return suggestion
 
-def extract_name_age_intelligent(text):
-    """Intelligently extract name and age from user input using multiple patterns."""
-    name, age = None, None
-    text_lower = text.lower().strip()
-    
-    # Age extraction patterns (more comprehensive)
-    age_patterns = [
-        r'(?:i am|i\'m|im|age is|aged|my age is|years old|year old)\s*(\d{1,3})',
-        r'(\d{1,3})\s*(?:years old|year old|yrs old|yr old)',
-        r'\b(\d{1,3})\s*(?:and|,)?\s*(?:years|yrs|y\.o\.)',
-        r'(?:^|\s)(\d{1,3})(?:\s|$)',  # standalone numbers
-    ]
-    
-    for pattern in age_patterns:
-        match = re.search(pattern, text_lower)
-        if match:
-            potential_age = int(match.group(1))
-            if 1 <= potential_age <= 120:  # reasonable age range
-                age = str(potential_age)
-                break
-    
-    # Name extraction patterns (more comprehensive)
-    name_patterns = [
-        r'(?:my name is|name is|i am|i\'m|im|call me|this is)\s+([a-zA-Z][a-zA-Z\s]{1,20}?)(?:\s+and|\s+\d|\s*$)',
-        r'^([a-zA-Z][a-zA-Z\s]{1,20}?)\s+(?:and|,)?\s*\d',  # name followed by number
-        r'(?:^|\s)([a-zA-Z]{2,15})(?:\s+and|\s+\d)',  # simple name pattern
-    ]
-    
-    for pattern in name_patterns:
-        match = re.search(pattern, text_lower)
-        if match:
-            potential_name = match.group(1).strip().title()
-            # Filter out common non-name words
-            non_names = ['it', 'is', 'am', 'my', 'me', 'the', 'and', 'or', 'but', 'yes', 'no']
-            if potential_name.lower() not in non_names and len(potential_name) >= 2:
-                name = potential_name
-                break
-    
-    # Special case: handle "thanush and 23" or "it thanush and im 23" patterns
-    special_patterns = [
-        r'(?:it\s+)?([a-zA-Z]{2,15})\s+and\s+(?:im\s+|i\'m\s+)?(\d{1,3})',
-        r'([a-zA-Z]{2,15})\s+and\s+(\d{1,3})',
-    ]
-    
-    for pattern in special_patterns:
-        match = re.search(pattern, text_lower)
-        if match:
-            potential_name = match.group(1).strip().title()
-            potential_age = int(match.group(2))
-            if potential_name.lower() not in ['it', 'is', 'am'] and 1 <= potential_age <= 120:
-                name = potential_name
-                age = str(potential_age)
-                break
-    
-    return name, age
-
-def extract_name_age_from_all_messages(messages):
-    """Extract name and age from all conversation messages."""
-    name, age = None, None
+def build_simple_prompt(system_prompt, conversation_history, current_input):
+    """Build a simple prompt for Llama-2"""
+    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
     
-    for msg in messages:
-        if msg.type == "human":
-            extracted_name, extracted_age = extract_name_age_intelligent(msg.content)
-            if extracted_name and not name:
-                name = extracted_name
-            if extracted_age and not age:
-                age = extracted_age
+    # Add conversation history
+    for i, (user_msg, bot_msg) in enumerate(conversation_history):
+        prompt += f"{user_msg} [/INST] {bot_msg} </s><s>[INST] "
     
-    return name, age
-
-def is_medical_symptom_message(text):
-    """Check if the message contains medical symptoms rather than just name/age."""
-    medical_keywords = [
-        'hurt', 'pain', 'ache', 'sick', 'fever', 'cough', 'headache', 'stomach', 'throat',
-        'nausea', 'dizzy', 'tired', 'fatigue', 'breathe', 'chest', 'back', 'leg', 'arm',
-        'symptom', 'feel', 'suffering', 'problem', 'issue', 'uncomfortable', 'sore'
-    ]
+    # Add current input
+    prompt += f"{current_input} [/INST] "
     
-    text_lower = text.lower()
-    return any(keyword in text_lower for keyword in medical_keywords)
+    return prompt
 
 @spaces.GPU
 def generate_response(message, history):
-    """Generate a response using both models, with full context."""
-    # Save the latest user message and last assistant response to memory
-    if history and len(history[-1]) == 2:
-        memory.save_context({"input": history[-1][0]}, {"output": history[-1][1]})
-    memory.save_context({"input": message}, {"output": ""})
-
-    messages = memory.chat_memory.messages
-    
-    # Extract name and age from all messages
-    name, age = extract_name_age_from_all_messages(messages)
-    
-    # Check what information is missing
-    missing_info = []
-    if not name:
-        missing_info.append("your name")
-    if not age:
-        missing_info.append("your age")
-    
-    # If missing basic info, ask for it
-    if missing_info:
-        ask = "Hello! Before we discuss your health concerns, could you please tell me " + " and ".join(missing_info) + "?"
-        return ask
-    
-    # Count meaningful medical information exchanges (exclude name/age only messages)
-    medical_info_turns = 0
-    for msg in messages:
-        if msg.type == "human":
-            # Count only if it's not just name/age info and contains medical content
-            if is_medical_symptom_message(msg.content) or not any(keyword in msg.content.lower() for keyword in ['name', 'age', 'years', 'old', 'im', 'i am']):
-                medical_info_turns += 1
+    """Generate a response using simple state tracking."""
+    global conversation_state
+    
+    # Reset state if this is a new conversation
+    if not history:
+        conversation_state = {
+            'name': None,
+            'age': None,
+            'medical_turns': 0,
+            'has_name': False,
+            'has_age': False
+        }
+    
+    # Step 1: Ask for name if not provided
+    if not conversation_state['has_name']:
+        conversation_state['has_name'] = True
+        return "Hello! Before we discuss your health concerns, could you please tell me your name?"
+    
+    # Step 2: Store name and ask for age
+    if conversation_state['name'] is None:
+        conversation_state['name'] = message.strip()
+        return f"Nice to meet you, {conversation_state['name']}! Could you please tell me your age?"
+    
+    # Step 3: Store age and start medical questions
+    if not conversation_state['has_age']:
+        conversation_state['age'] = message.strip()
+        conversation_state['has_age'] = True
+        return f"Thank you, {conversation_state['name']}! Now, what brings you here today? Please tell me about any symptoms or health concerns you're experiencing."
+    
+    # Step 4: Medical consultation phase
+    conversation_state['medical_turns'] += 1
+    
+    # Prepare conversation history for the model
+    medical_history = []
+    if len(history) >= 3:  # Skip name/age exchanges
+        medical_history = history[3:]
+    
+    # Define follow-up questions based on turn number
+    followup_questions = [
+        "Can you describe your symptoms in more detail? What exactly are you experiencing?",
+        "How long have you been experiencing these symptoms? When did they first start?",
+        "On a scale of 1-10, how would you rate the severity of your symptoms?",
+        "Have you noticed anything that makes your symptoms better or worse?",
+        "Do you have any other symptoms, medical history, or are you taking any medications?"
+    ]
     
-    # Ensure we have at least one medical symptom mentioned
-    if medical_info_turns == 0 and not is_medical_symptom_message(message):
-        return f"Thank you, {name}! Now, what brings you here today? Please tell me about any symptoms or health concerns you're experiencing."
+    # Build the prompt for medical consultation
+    if conversation_state['medical_turns'] <= 5:
+        # Still gathering information
+        prompt = build_simple_prompt(SYSTEM_PROMPT, medical_history, message)
+        
+        # Generate response
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = 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=tokenizer.eos_token_id
+            )
+        
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        llama_response = full_response.split('[/INST]')[-1].strip()
+        
+        # Add a specific follow-up question
+        if conversation_state['medical_turns'] < len(followup_questions):
+            next_question = followup_questions[conversation_state['medical_turns']]
+            return f"{llama_response}\n\n{next_question}"
+        else:
+            return llama_response
     
-    # Ask up to 5 intelligent follow-up questions, then provide diagnosis and treatment
-    if medical_info_turns < 5:
-        prompt = build_llama2_prompt(SYSTEM_PROMPT, messages, message, followup_stage=medical_info_turns)
     else:
-        # Time for final diagnosis and treatment recommendations
-        prompt = build_llama2_prompt(SYSTEM_PROMPT, messages, message)
-        prompt = prompt.replace("[/INST] ", "[/INST] Based on all the information provided, please provide a comprehensive assessment including: 1) Most likely diagnosis, 2) Recommended next steps, and 3) When to seek immediate medical attention. ")
-
-    # Generate response using Llama-2
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    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.eos_token_id
-        )
-    
-    full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
-    llama_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
-
-    # After 5 medical info turns, add Meditron suggestions
-    if medical_info_turns >= 4:  # Start suggesting after 4+ turns
-        # Compile patient information for Meditron
-        patient_summary = f"Patient: {name}, Age: {age}\n\n"
-        patient_summary += "Medical Information:\n"
+        # Time for diagnosis and treatment (after 5+ turns)
+        # Compile patient information
+        patient_info = f"Patient: {conversation_state['name']}, Age: {conversation_state['age']}\n\n"
+        patient_info += "Symptoms and Information:\n"
         
-        for msg in messages:
-            if msg.type == "human" and is_medical_symptom_message(msg.content):
-                patient_summary += f"- {msg.content}\n"
+        # Add all medical conversation history
+        for user_msg, bot_msg in medical_history:
+            patient_info += f"Patient: {user_msg}\n"
+        patient_info += f"Patient: {message}\n"
         
-        patient_summary += f"\nLatest input: {message}\n"
-        patient_summary += f"\nInitial Assessment: {llama_response}"
+        # Generate diagnosis with Llama-2
+        diagnosis_prompt = f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\nBased on all the information provided, please provide a comprehensive medical assessment including likely diagnosis and recommendations for {conversation_state['name']}.\n\nPatient Information:\n{patient_info} [/INST] "
         
-        # Get Meditron suggestions
-        medicine_suggestions = get_meditron_suggestions(patient_summary)
+        inputs = tokenizer(diagnosis_prompt, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                max_new_tokens=384,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        diagnosis = full_response.split('[/INST]')[-1].strip()
+        
+        # Get treatment suggestions from Meditron
+        treatment_suggestions = get_meditron_suggestions(patient_info)
+        
+        # Combine responses
+        final_response = f"{diagnosis}\n\n--- TREATMENT RECOMMENDATIONS ---\n\n{treatment_suggestions}\n\n**Important:** These are general recommendations. Please consult with a healthcare professional for personalized medical advice."
         
-        final_response = (
-            f"{llama_response}\n\n"
-            f"--- MEDICATION AND HOME CARE RECOMMENDATIONS ---\n\n"
-            f"{medicine_suggestions}\n\n"
-            f"**Important:** These are general recommendations. Please consult with a healthcare professional for personalized medical advice, especially if symptoms persist or worsen."
-        )
         return final_response
 
-    return llama_response
-
 # Create the Gradio interface
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="🩺 AI Medical Assistant with Treatment Suggestions",
-    description="Describe your symptoms and I'll gather information to provide medical insights and treatment recommendations.",
+    title="🩺 AI Medical Assistant",
+    description="I'll ask for your basic information first, then gather details about your symptoms to provide medical insights.",
     examples=[
-        "Hi, I'm Sarah and I'm 25. I have a persistent cough and sore throat.",
-        "My name is John, I'm 35, and I've been having severe headaches.",
-        "I'm Lisa, 28 years old, and my stomach has been hurting since yesterday."
+        "I have a persistent cough",
+        "I've been having headaches",
+        "My stomach hurts"
     ],
     theme="soft"
 )