Made changes to chatbot.py

app/chatbot.py

@@ -76,6 +76,31 @@ def is_affirmative(answer):
     answer_lower = answer.lower()
     return any(word in answer_lower for word in ["yes", "yeah", "yep", "certainly", "sometimes", "a little"])
 import random
+# Map symptoms to associated body parts
+SYMPTOM_BODY_PARTS = {
+    "headache": "head",
+    "fever": None,
+    "cold": None,
+    "cough": "throat",
+    "chest pain": "chest",
+    "stomach ache": "abdomen",
+    "sore throat": "throat",
+    "back pain": "back",
+    "ear pain": "ear",
+    "toothache": "mouth",
+    "nausea": "stomach",
+    "vomiting": "stomach",
+    "diarrhea": "abdomen",
+    "shortness of breath": "chest",
+    "dizziness": "head",
+    "rash": "skin",
+    "eye pain": "eye",
+    "abdominal pain": "abdomen",
+    "joint pain": "joint",
+    "muscle pain": "muscle",
+    "neck pain": "neck"
+}
+
 # -------------------------------
 # Chatbot session class
 # -------------------------------
@@ -132,7 +157,7 @@ class ChatbotSession:
             if unexplored_symptoms:
                 symptom = random.choice(unexplored_symptoms)
                 self.current_detail_symptom = symptom
-                self.symptom_details[symptom] = {}  # Start collecting details only when asked
+                self.symptom_details[symptom] = {}
                 self.state = "symptom_detail"
                 return f"Doctor: About your '{symptom}', when did it start? (e.g., 2 days ago)"
             else:
@@ -149,7 +174,13 @@ class ChatbotSession:
 
         if self.current_detail_symptom and 'severity' not in self.symptom_details[self.current_detail_symptom]:
             self.symptom_details[self.current_detail_symptom]['severity'] = message
-            return f"Doctor: Where exactly do you feel the '{self.current_detail_symptom}'? (e.g., forehead, chest, abdomen)"
+            body_part = SYMPTOM_BODY_PARTS.get(self.current_detail_symptom)
+            if body_part:
+                self.symptom_details[self.current_detail_symptom]['location'] = body_part
+                self.state = "intake"
+                return "Doctor: Thank you. Are there any more symptoms you'd like to mention?"
+            else:
+                return f"Doctor: Where exactly do you feel the '{self.current_detail_symptom}'? (e.g., forehead, chest, abdomen)"
 
         if self.current_detail_symptom and 'location' not in self.symptom_details[self.current_detail_symptom]:
             self.symptom_details[self.current_detail_symptom]['location'] = message
@@ -182,11 +213,11 @@ class ChatbotSession:
             matches = len(set(self.reported_symptoms) & set(symptoms))
             if matches > 0:
                 score = matches / len(symptoms)
-                if matches >= 2:  # Only consider diseases with at least 2 matching symptoms
+                if matches >= 2:
                     match_scores.append((disease, score))
 
         match_scores.sort(key=lambda x: x[1], reverse=True)
-        self.predicted_diseases = match_scores[:2]  # Top 2 predictions
+        self.predicted_diseases = match_scores[:2]
 
     def _generate_summary(self):
         report_lines = []