Made changes to chatbot.py

app/chatbot.py

@@ -39,11 +39,20 @@ index.add(np.array(tfidf_matrix, dtype=np.float32))
 disease_list = list(disease_symptoms.keys())
 
 def find_closest_disease(user_symptoms):
-    if not user_symptoms:
+    max_matches = 0
+    best_disease = None
+
+    for disease, symptoms in disease_symptoms.items():
+        matches = len(set(user_symptoms) & set(symptoms))
+        if matches > max_matches:
+            max_matches = matches
+            best_disease = disease
+
+    # Require at least 2 symptoms matching to consider it valid
+    if max_matches >= 2:
+        return best_disease
+    else:
         return None
-    user_vector = vectorizer.transform([" ".join(user_symptoms)]).toarray().astype("float32")
-    distances, indices = index.search(user_vector, k=1)
-    return disease_list[indices[0][0]]
 
 # -------------------------------
 # Load Medical NER model for symptom extraction
@@ -74,23 +83,22 @@ class ChatbotSession:
     def __init__(self):
         self.conversation_history = []
         self.reported_symptoms = set()
-        self.symptom_details = {}  # New: Track details per symptom
+        self.symptom_details = {}
         self.asked_missing = set()
         self.awaiting_followup = None
-        self.awaiting_detail = None
         self.current_detail_symptom = None
-        self.state = "symptom_collection"
+        self.state = "intake"
         self.finished = False
-        self.predicted_disease = None
-        greeting = "Doctor: Hello, I am your virtual doctor. What symptoms are you experiencing today?"
+        self.predicted_diseases = []
+        greeting = "Doctor: Hello! I'm here to assist you. Could you please describe how you're feeling today?"
         self.conversation_history.append(greeting)
 
     def process_message(self, message: str) -> str:
         if self.finished:
-            return "Doctor: Thank you. Our session has ended."
+            return "Doctor: Thank you. Our session has ended. Wishing you good health!"
 
-        if self.state == "symptom_collection":
-            return self._handle_symptom_collection(message)
+        if self.state == "intake":
+            return self._handle_intake(message)
 
         if self.state == "symptom_detail":
             return self._handle_symptom_detail(message)
@@ -101,20 +109,19 @@ class ChatbotSession:
         if self.state == "medications":
             return self._handle_medications(message)
 
-        return "Doctor: Could you please clarify?"
+        return "Doctor: Could you please clarify that?"
 
-    def _handle_symptom_collection(self, message):
+    def _handle_intake(self, message):
         if message.lower() in ["exit", "quit", "no"]:
             if not self.reported_symptoms:
                 goodbye = "Doctor: It seems no symptoms were reported. Ending the session."
                 self.finished = True
                 return goodbye
             else:
-                self.predicted_disease = find_closest_disease(list(self.reported_symptoms))
+                self._predict_diseases()
                 self.state = "pain_check"
-                return f"Doctor: Before proceeding, are you experiencing any pain? If yes, please rate it 1-10 or type 'no'."
+                return "Doctor: Before proceeding, are you experiencing any pain? If yes, please rate it 1-10 or type 'no'."
 
-        # Extract symptoms
         ner_results = extract_symptoms_ner(message)
         if ner_results:
             for sym in ner_results:
@@ -122,52 +129,44 @@ class ChatbotSession:
                     self.reported_symptoms.add(sym)
                     self.symptom_details[sym] = {}
 
-            # Update disease prediction
-            self.predicted_disease = find_closest_disease(list(self.reported_symptoms))
-            if self.predicted_disease:
-                expected = set(disease_symptoms.get(self.predicted_disease, []))
-                missing = expected - self.reported_symptoms
-                not_asked = missing - self.asked_missing
+            self._predict_diseases()
 
-                if not_asked:
-                    symptom_to_ask = list(not_asked)[0]
-                    self.awaiting_followup = symptom_to_ask
-                    return f"Doctor: Are you also experiencing {symptom_to_ask}?"
-
-            # If all covered, ask symptom details
             if self.reported_symptoms:
-                symptom = random.choice(list(self.reported_symptoms))
+                symptom = random.choice(list(self.reported_symptoms - set(self.symptom_details.keys())))
                 self.current_detail_symptom = symptom
                 self.state = "symptom_detail"
-                return f"Doctor: About your '{symptom}', when did it start? (duration)"
+                return f"Doctor: About your '{symptom}', when did it start? (e.g., 2 days ago)"
 
         else:
-            return "Doctor: I couldn't detect any medical symptoms. Could you describe it differently?"
+            return "Doctor: I couldn't clearly detect any medical symptoms. Could you please describe your condition differently?"
 
-        return "Doctor: Do you have any more symptoms to share?"
+        return "Doctor: Any other symptoms you'd like to share?"
 
     def _handle_symptom_detail(self, message):
         if self.current_detail_symptom and 'duration' not in self.symptom_details[self.current_detail_symptom]:
             self.symptom_details[self.current_detail_symptom]['duration'] = message
-            return f"Doctor: How severe is the '{self.current_detail_symptom}'? (mild/moderate/severe)"
+            return f"Doctor: How severe is the '{self.current_detail_symptom}'? (mild, moderate, severe)"
 
         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}' (body part/location)?"
+            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
-            self.state = "symptom_collection"
-            return "Doctor: Thank you. Any other symptoms you'd like to mention?"
+            self.state = "intake"
+            return "Doctor: Thank you. Are there any more symptoms you'd like to mention?"
 
-        return "Doctor: Please clarify."
+        return "Doctor: Please clarify that."
 
     def _handle_pain_check(self, message):
         try:
             pain_level = int(message)
             self.symptom_details['pain'] = {'severity': pain_level}
         except ValueError:
-            self.symptom_details['pain'] = {'severity': message}
+            if is_affirmative(message):
+                self.symptom_details['pain'] = {'severity': 'mild'}
+            else:
+                self.symptom_details['pain'] = {'severity': 'none'}
 
         self.state = "medications"
         return "Doctor: Have you taken any medications recently? Please mention them or type 'no'."
@@ -175,9 +174,18 @@ class ChatbotSession:
     def _handle_medications(self, message):
         self.symptom_details['medications'] = message if message.lower() not in ["no", "none"] else "None"
         self.finished = True
+        return self._generate_summary()
+
+    def _predict_diseases(self):
+        match_scores = []
+        for disease, symptoms in disease_symptoms.items():
+            matches = len(set(self.reported_symptoms) & set(symptoms))
+            if matches > 0:
+                score = matches / len(symptoms)
+                match_scores.append((disease, score))
 
-        summary = self._generate_summary()
-        return summary
+        match_scores.sort(key=lambda x: x[1], reverse=True)
+        self.predicted_diseases = match_scores[:3]  # Top 3 predictions
 
     def _generate_summary(self):
         report = "\n".join([
@@ -185,15 +193,22 @@ class ChatbotSession:
             for sym, details in self.symptom_details.items()
         ])
 
-        disease_part = f"\n\nDoctor: Based on your symptoms, the most likely condition could be *{self.predicted_disease}*." if self.predicted_disease else ""
-        advice = "\n\nDoctor: Please note this is a preliminary assessment. Kindly consult a physician for a detailed diagnosis."
-        
-        return f"Doctor: Thank you for the information! Here's a summary of your case:\n{report}{disease_part}{advice}"
+        disease_part = ""
+        if self.predicted_diseases:
+            disease_part = "\n\nPossible Conditions:" + "\n".join([
+                f"- {disease} ({int(score * 100)}%)" for disease, score in self.predicted_diseases
+            ])
+        else:
+            disease_part = "\n\nDoctor: No strong condition match was found based on reported symptoms."
+
+        advice = "\n\nDoctor: Please note this is a preliminary virtual consultation. You should consult a physician if symptoms persist, worsen, or if you feel uncomfortable."
+
+        return f"Doctor: Thank you for the detailed information!\n\n=== Medical Report ===\n{report}{disease_part}{advice}"
 
     def get_data(self):
         return {
             "conversation": self.conversation_history,
             "symptoms": list(self.reported_symptoms),
             "symptom_details": self.symptom_details,
-            "predicted_disease": self.predicted_disease
+            "predicted_diseases": self.predicted_diseases
         }