Update app.py

app.py

@@ -47,20 +47,20 @@ def extract_name(text):
 # =======================================================
 def doctor_response(user_message):
     global session
-    user_message = user_message.strip()
+    user_message = user_message.strip().lower()
 
-    # Intro
+    # Step 1: Greeting
     if session["stage"] == "intro":
         session["stage"] = "ask_name"
         return "👨‍⚕️ Hello! I’m Dr. Aiden. May I know your name, please?"
 
-    # Ask Name
+    # Step 2: Get Name
     elif session["stage"] == "ask_name":
         session["name"] = extract_name(user_message)
         session["stage"] = "ask_age"
         return f"Nice to meet you, {session['name']}! How old are you?"
 
-    # Ask Age
+    # Step 3: Get Age
     elif session["stage"] == "ask_age":
         words = user_message.split()
         for w in words:
@@ -70,54 +70,102 @@ def doctor_response(user_message):
                 return f"Got it, {session['name']}. Are you male or female?"
         return "Please tell me your age in numbers, like 20 or 25."
 
-    # Ask Gender
+    # Step 4: Get Gender
     elif session["stage"] == "ask_gender":
-        if "male" in user_message.lower():
+        if "male" in user_message:
             session["gender"] = "male"
-        elif "female" in user_message.lower():
+        elif "female" in user_message:
             session["gender"] = "female"
         else:
             return "Could you please specify whether you are male or female?"
         session["stage"] = "ask_symptoms"
         return f"Thanks, {session['name']}! So you're a {session['age']}-year-old {session['gender']}. What symptoms are you experiencing?"
 
-    # Ask Symptoms
+    # Step 5: Get Symptoms
     elif session["stage"] == "ask_symptoms":
         session["symptoms"] = user_message
         session["stage"] = "ask_duration"
         return "Since when have you been feeling this way?"
 
-    # Ask Duration
+    # Step 6: Duration
     elif session["stage"] == "ask_duration":
         session["duration"] = user_message
-        session["stage"] = "consult"
+        session["stage"] = "ask_medicine"
         return "Got it. Are you taking any medications or treatments currently?"
 
-    # Consultation
+    # Step 7: Medicine intake
+    elif session["stage"] == "ask_medicine":
+        session["medication"] = user_message
+        session["stage"] = "consult"
+        return "Thank you. Let’s discuss what could be going on and how you can manage it safely."
+
+    # Step 8: Personalized Consultation
     elif session["stage"] == "consult":
         name = session["name"]
         age = session["age"]
         gender = session["gender"]
-        symptoms = session["symptoms"]
-        duration = session["duration"]
-
+        symptoms = session.get("symptoms", "")
+        duration = session.get("duration", "")
+        medication = session.get("medication", "")
+
+        # ------------- Handle diet-related questions -----------------
+        if any(word in user_message for word in ["diet", "food", "meal", "eat", "nutrition"]):
+            # Diet suggestions tailored to the condition
+            if "fever" in symptoms:
+                diet = (
+                    f"🍎 {name}, since you have a fever, try keeping your meals light and hydrating.\n\n"
+                    "- 🥣 **Breakfast:** Oatmeal or boiled egg with fruit (like banana or apple)\n"
+                    "- 🍲 **Lunch:** Rice with lentil soup or boiled vegetables\n"
+                    "- 🍛 **Dinner:** Light soup or grilled chicken with plain rice\n"
+                    "- 💧 **Hydration:** Drink 8–10 glasses of water, coconut water, or clear soups\n"
+                    "- 🚫 **Avoid:** Fried, oily, or spicy foods\n\n"
+                    "Eat small portions often — this helps your body recover faster."
+                )
+            elif "stomach" in symptoms or "vomit" in symptoms:
+                diet = (
+                    f"🥗 {name}, for stomach discomfort, stick to a **bland diet**:\n\n"
+                    "- 🍞 Toast, plain rice, boiled potatoes, or bananas\n"
+                    "- 💧 Sip water, oral rehydration solution, or herbal tea\n"
+                    "- 🚫 Avoid milk, fried, or spicy foods\n"
+                    "- 🍌 Eat small meals to avoid nausea"
+                )
+            elif "cold" in symptoms or "flu" in symptoms or "cough" in symptoms:
+                diet = (
+                    f"🍊 {name}, for flu or cold, focus on **immunity-boosting foods**:\n\n"
+                    "- 🍋 Citrus fruits, honey with warm water, and soups\n"
+                    "- 🍲 Chicken soup helps clear congestion\n"
+                    "- ☕ Ginger or green tea for throat relief\n"
+                    "- 🚫 Avoid sugary and chilled drinks"
+                )
+            else:
+                diet = (
+                    f"🥗 {name}, eat a balanced diet: fruits, vegetables, lean protein, and whole grains. "
+                    "Stay hydrated and avoid processed or fried foods."
+                )
+
+            return diet + "\n\n⚕️ *Note: This is general advice — not a substitute for medical care.*"
+
+        # ----------------- Normal medical consultation ----------------
         prompt = f"""
-You are Dr. Aiden — a warm, friendly, and professional doctor having an interview-style consultation.
-The patient is a {age}-year-old {gender} named {name}.
-They have been feeling {symptoms} for {duration}.
-They said: "{user_message}"
+You are Dr. Aiden — a warm, caring, and professional doctor.
+You are consulting a {age}-year-old {gender} named {name}.
 
-Respond like a real doctor — show empathy, analyze the symptoms, suggest likely causes, give simple medication and home care advice.
+Patient details:
+- Symptoms: {symptoms}
+- Duration: {duration}
+- Current medications: {medication}
 
-Include:
-1. Acknowledge their discomfort.
-2. Explain possible causes in simple terms.
-3. Recommend over-the-counter medicines (if safe).
-4. Suggest food, hydration, and rest tips.
-5. Warn when to visit a real doctor.
-6. End with gentle reassurance.
+Now the patient says: "{user_message}"
 
-Doctor:"""
+Respond as a real doctor would — empathetic, clear, and personalized.
+Include:
+1. Acknowledge their condition
+2. Possible causes
+3. Simple home remedies or OTC medicines (if safe)
+4. Diet, rest, and hydration tips
+5. When to visit a real doctor
+6. End with a reassuring tone
+"""
 
         inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
         gen_cfg = GenerationConfig(
@@ -134,13 +182,15 @@ Doctor:"""
 
         output_text = tokenizer.decode(output[0], skip_special_tokens=True)
         output_text = output_text.split("Doctor:")[-1].strip()
+
         if not output_text.endswith((".", "!", "?")):
             output_text += "."
-        output_text
 
+        output_text += "\n\n⚕️ *Note: This advice is AI-generated and not a substitute for professional medical care.*"
         return output_text
 
 
+
 # =======================================================
 # Gradio Interface
 # =======================================================