fix

app.py

@@ -114,6 +114,7 @@ def chat_with_model(messages):
         # Early stopping if user reappears
         if "\nUser:" in output_text:
             output_text = output_text.split("\nUser:")[0].rstrip()
+            messages[-1]["content"] = output_text
             break
 
         # Strip prompt from start of generated output
@@ -135,8 +136,7 @@ def chat_with_model(messages):
     if in_think:
         output_text += "*"
         messages[-1]["content"] = output_text
-        yield messages
-
+    
     # Wait for thread to finish
     thread.join(timeout=1.0)
     current_model.to("cpu")
@@ -145,7 +145,7 @@ def chat_with_model(messages):
     messages[-1]["content"] = output_text
     print(f'Step 3: {messages}')
 
-    return messages
+    yield messages