Enhance user information collection in app.py by emphasizing follow-up questions and refining response generation logic based on actual information turns.

app.py

@@ -13,7 +13,7 @@ SYSTEM_PROMPT = """You are a professional virtual doctor. Your goal is to collec
 
 Always begin by asking for the user's name and age if not already provided.
 
-Ask 1-2 follow-up questions at a time to gather more details about:
+**IMPORTANT** Ask 1-2 follow-up questions at a time to gather more details about:
 - Detailed description of symptoms
 - Duration (when did it start?)
 - Severity (scale of 1-10)
@@ -99,10 +99,10 @@ def extract_name_age(messages):
     for msg in messages:
         if msg.type == "human":
             age_match = re.search(r"(?:I am|I'm|age is|aged|My age is)\s*(\d{1,3})", msg.content, re.IGNORECASE)
-            if age_match:
+            if age_match and not age:
                 age = age_match.group(1)
             name_match = re.search(r"(?:my name is|I'm|I am)\s*([A-Za-z]+)", msg.content, re.IGNORECASE)
-            if name_match:
+            if name_match and not name:
                 name = name_match.group(1)
     return name, age
 
@@ -125,10 +125,17 @@ def generate_response(message, history):
         ask = "Before we continue, 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
+    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
+
     prompt = build_llama2_prompt(SYSTEM_PROMPT, messages, message)
-    num_user_turns = sum(1 for m in messages if m.type == "human")
-    # Only add summarization ONCE, not on every turn after 4
-    if num_user_turns == 4:
+    # Only add summarization ONCE, not on every turn after 4 info turns
+    if info_turns == 4:
         prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -149,9 +156,11 @@ def generate_response(message, history):
     full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
     llama_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
 
-    # After 4 turns, add medicine suggestions from Meditron, but only once
-    if num_user_turns == 4:
-        full_patient_info = "\n".join([m.content for m in messages if m.type == "human"] + [message]) + "\n\nSummary: " + llama_response
+    # After 4 info turns, add medicine suggestions from Meditron, but only once
+    if info_turns == 4:
+        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)
         final_response = (
             f"{llama_response}\n\n"