Refactor app.py to extract user name and age from conversation history and improve response generation logic

app.py

@@ -3,6 +3,7 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from langchain.memory import ConversationBufferMemory
+import re
 
 # Model configuration
 LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
@@ -93,6 +94,18 @@ 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):
+    name, age = None, None
+    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:
+                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:
+                name = name_match.group(1)
+    return name, age
+
 @spaces.GPU
 def generate_response(message, history):
     """Generate a response using both models, with full context."""
@@ -101,15 +114,21 @@ def generate_response(message, history):
         memory.save_context({"input": history[-1][0]}, {"output": history[-1][1]})
     memory.save_context({"input": message}, {"output": ""})
 
-    # Use the full message sequence from memory
     messages = memory.chat_memory.messages
+    name, age = extract_name_age(messages)
+    missing_info = []
+    if not name:
+        missing_info.append("your name")
+    if not age:
+        missing_info.append("your age")
+    if missing_info:
+        ask = "Before we continue, could you please tell me " + " and ".join(missing_info) + "?"
+        return ask
 
-    # Build the prompt with the full message sequence
     prompt = build_llama2_prompt(SYSTEM_PROMPT, messages, message)
-
-    # Add summarization instruction after 4 turns (count human messages)
     num_user_turns = sum(1 for m in messages if m.type == "human")
-    if num_user_turns >= 4:
+    # Only add summarization ONCE, not on every turn after 4
+    if num_user_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)
@@ -130,15 +149,10 @@ 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
-    if num_user_turns >= 4:
-        # Collect full patient conversation (all user messages)
+    # 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
-
-        # Get medicine suggestions
         medicine_suggestions = get_meditron_suggestions(full_patient_info)
-
-        # Format final response
         final_response = (
             f"{llama_response}\n\n"
             f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"