Enhance user interaction in app.py by refining follow-up questions for symptom collection and implementing intelligent extraction of user name and age from messages. Improve response generation logic to ensure comprehensive medical assessments and treatment recommendations.

app.py

@@ -77,14 +77,15 @@ def build_llama2_prompt(system_prompt, messages, user_input, followup_stage=None
             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 detail?",
-            "How long have you been experiencing these symptoms?",
-            "On a scale of 1-10, how severe are your symptoms?",
-            "Have you noticed anything that makes your symptoms better or worse?",
-            "Do you have any other related symptoms, such as fever, fatigue, or shortness of breath?"
+            "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] "
@@ -112,27 +113,87 @@ 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(messages):
+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
+    
     for msg in messages:
         if msg.type == "human":
-            # Try to extract age
-            age_match = re.search(r"(?:I am|I'm|age is|aged|My age is|im|i'm)\s*(\d{1,3})", msg.content, re.IGNORECASE)
-            if age_match and not age:
-                age = age_match.group(1)
-            # Try to extract name (avoid matching 'I'm' as name if age is present)
-            name_match = re.search(r"my name is\s*([A-Za-z]+)", msg.content, re.IGNORECASE)
-            if name_match and not name:
-                name = name_match.group(1)
-            # Fallback: if user says 'I'm <name> and <age>'
-            fallback_match = re.search(r"i['’`]?m\s*([A-Za-z]+)\s*(?:and|,)?\s*(\d{1,3})", msg.content, re.IGNORECASE)
-            if fallback_match:
-                if not name:
-                    name = fallback_match.group(1)
-                if not age:
-                    age = fallback_match.group(2)
+            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
+    
     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'
+    ]
+    
+    text_lower = text.lower()
+    return any(keyword in text_lower for keyword in medical_keywords)
+
 @spaces.GPU
 def generate_response(message, history):
     """Generate a response using both models, with full context."""
@@ -142,31 +203,43 @@ def generate_response(message, history):
     memory.save_context({"input": message}, {"output": ""})
 
     messages = memory.chat_memory.messages
-    name, age = extract_name_age(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 = "Before we continue, could you please tell me " + " and ".join(missing_info) + "?"
+        ask = "Hello! Before we discuss your health concerns, could you please tell me " + " and ".join(missing_info) + "?"
         return ask
-
-    # Count how many user turns have actually provided new info (not just name/age)
-    info_turns = 0
+    
+    # Count meaningful medical information exchanges (exclude name/age only messages)
+    medical_info_turns = 0
     for msg in messages:
         if msg.type == "human":
-            # Ignore turns that only provide name/age
-            if not re.fullmatch(r".*(name|age|years? old|I'm|I am|my name is).*", msg.content, re.IGNORECASE):
-                info_turns += 1
-
-    # Ask up to 5 intelligent follow-up questions, then summarize/diagnose
-    if info_turns < 5:
-        prompt = build_llama2_prompt(SYSTEM_PROMPT, messages, message, followup_stage=info_turns)
+            # 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
+    
+    # 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."
+    
+    # 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] Now, based on all the information, provide a likely diagnosis (if possible), and suggest when professional care may be needed. ")
+        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(
@@ -178,19 +251,31 @@ def generate_response(message, history):
             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 info turns, add medicine suggestions from Meditron, but only once
-    if info_turns == 5:
-        full_patient_info = "\n".join([
-            m.content for m in messages if m.type == "human" and not re.fullmatch(r".*(name|age|years? old|I'm|I am|my name is).*", m.content, re.IGNORECASE)
-        ] + [message]) + "\n\nSummary: " + llama_response
-        medicine_suggestions = get_meditron_suggestions(full_patient_info)
+    # 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"
+        
+        for msg in messages:
+            if msg.type == "human" and is_medical_symptom_message(msg.content):
+                patient_summary += f"- {msg.content}\n"
+        
+        patient_summary += f"\nLatest input: {message}\n"
+        patient_summary += f"\nInitial Assessment: {llama_response}"
+        
+        # Get Meditron suggestions
+        medicine_suggestions = get_meditron_suggestions(patient_summary)
+        
         final_response = (
             f"{llama_response}\n\n"
-            f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
-            f"{medicine_suggestions}"
+            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
 
@@ -199,12 +284,12 @@ def generate_response(message, history):
 # Create the Gradio interface
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Medical Assistant with Medicine Suggestions",
-    description="Tell me about your symptoms, and after gathering enough information, I'll suggest potential remedies.",
+    title="🩺 AI Medical Assistant with Treatment Suggestions",
+    description="Describe your symptoms and I'll gather information to provide medical insights and treatment recommendations.",
     examples=[
-        "I have a cough and my throat hurts",
-        "I've been having headaches for a week",
-        "My stomach has been hurting since yesterday"
+        "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."
     ],
     theme="soft"
 )