Fix initial prompt injection

app.py

@@ -97,17 +97,20 @@ def chat_with_model(messages):
         yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
         return
 
+    current_id = patient_id.value
+    if current_id is None:
+        yield messages
+        return
+
     pad_id = current_tokenizer.pad_token_id
     eos_id = current_tokenizer.eos_token_id
     if pad_id is None:
         pad_id = current_tokenizer.unk_token_id or 0
 
-    output_text = ""
-    in_think = False
-    max_new_tokens = 1024
-    generated_tokens = 0
+    # Remove the initial welcome if present
+    filtered_messages = [msg for msg in messages if not (msg["role"] == "assistant" and "Welcome to the Radiologist's Companion" in msg["content"])]
 
-    # ⮕ PREPARE SYSTEM + INITIAL MESSAGES
+    # Build system context
     system_messages = [
         {
             "role": "system",
@@ -132,26 +135,8 @@ def chat_with_model(messages):
         }
     ]
 
-    # Optional: if you later add available_images, you could append another system message.
-
-    welcome_message = (
-        "**Welcome to the Radiologist's Companion!**\n\n"
-        "You can ask me about the patient's medical history or available imaging data.\n"
-        "- I can summarize key details from the EHR.\n"
-        "- I can tell you which medical images are available.\n"
-        "- If you'd like an organ segmentation (e.g. spleen, liver, kidney_left, colon, femur_right) on an abdominal CT scan, just ask!\n\n"
-        "**Example Requests:**\n"
-        "- \"What do we know about this patient?\"\n"
-        "- \"Which images are available for this patient?\"\n"
-        "- \"Can you segment the spleen from the CT scan?\"\n"
-    )
-
-    # If it's the first user message (i.e., no assistant yet), prepend welcome
-    if len(messages) == 1 and messages[0]['role'] == 'user':
-        messages = [{"role": "assistant", "content": welcome_message}] + messages
+    full_messages = system_messages + filtered_messages
 
-    # Merge full conversation
-    full_messages = system_messages + messages
 
     prompt = format_prompt(full_messages)
 
@@ -277,31 +262,19 @@ def add_user_message(user_input, history):
     return "", patient_conversations[current_id]
 
 
-
 def autofill_patient(patient_key):
     if patient_key in patient_db:
         info = patient_db[patient_key]
 
-        # Init conversation if not existing
+        # Init empty conversation if not existing
         if info["id"] not in patient_conversations:
-            welcome_message = (
-                "**Welcome to the Radiologist's Companion!**\n\n"
-                "You can ask me about the patient's medical history or available imaging data.\n"
-                "- I can summarize key details from the EHR.\n"
-                "- I can tell you which medical images are available.\n"
-                "- If you'd like an organ segmentation (e.g. spleen, liver, kidney_left, colon, femur_right) on an abdominal CT scan, just ask!\n\n"
-                "**Example Requests:**\n"
-                "- \"What do we know about this patient?\"\n"
-                "- \"Which images are available for this patient?\"\n"
-                "- \"Can you segment the spleen from the CT scan?\"\n"
-            )
-
-            patient_conversations[info["id"]] = [{"role": "assistant", "content": welcome_message}]
+            patient_conversations[info["id"]] = []
 
         return info["name"], info["age"], info["id"], info["notes"]
     return "", "", "", ""
 
 
+
 with gr.Blocks(css=".gradio-container {height: 100vh; overflow: hidden;}") as demo:
     gr.Markdown("<h2 style='text-align: center;'>Radiologist's Companion</h2>")
 
@@ -347,11 +320,28 @@ with gr.Blocks(css=".gradio-container {height: 100vh; overflow: hidden;}") as de
     # After patient selected, load their conversation into chatbot
     def load_patient_conversation(patient_key):
         if patient_key in patient_db:
-            patient_id = patient_db[patient_key]["id"]
-            history = patient_conversations.get(patient_id, [])
-            return history
+            patient_id_val = patient_db[patient_key]["id"]
+            history = patient_conversations.get(patient_id_val, [])
+
+            # Show welcome + history
+            welcome_message = {
+                "role": "assistant",
+                "content": (
+                    "**Welcome to the Radiologist's Companion!**\n\n"
+                    "You can ask me about the patient's medical history or available imaging data.\n"
+                    "- I can summarize key details from the EHR.\n"
+                    "- I can tell you which medical images are available.\n"
+                    "- If you'd like an organ segmentation (e.g. spleen, liver, kidney_left, colon, femur_right) on an abdominal CT scan, just ask!\n\n"
+                    "**Example Requests:**\n"
+                    "- \"What do we know about this patient?\"\n"
+                    "- \"Which images are available for this patient?\"\n"
+                    "- \"Can you segment the spleen from the CT scan?\"\n"
+                )
+            }
+            return [welcome_message] + history
         return []
 
+
     patient_selector.change(
         autofill_patient,
         inputs=[patient_selector],