From 3a3e247f3cb21ec7c13bcfa9f21757e216d5c7ec Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 30 Aug 2025 12:36:35 -0700
Subject: [PATCH] Even better way to handle continue for thinking blocks

---
 modules/chat.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/modules/chat.py b/modules/chat.py
index 6d85bc6e..ad2f4001 100644
--- a/modules/chat.py
+++ b/modules/chat.py
@@ -333,20 +333,19 @@ def generate_chat_prompt(user_input, state, **kwargs):
 
         if _continue:
             prompt = prompt.split("fake assistant message replace me", 1)[0]
-            content = last_message.get("content", "")
-            thinking = last_message.get("thinking", "")
-            reasoning = last_message.get("reasoning_content", "")
 
-            partial_thought = thinking or reasoning
+            content = last_message.get("content", "")
+            partial_thought = last_message.get("thinking", "") or last_message.get("reasoning_content", "")
+
             # Handle partial thinking blocks (GPT-OSS and Seed-OSS)
-            if partial_thought and partial_thought.strip():
+            if not content and partial_thought and partial_thought.strip():
                 search_string = partial_thought.strip()
                 index = prompt.rfind(search_string)
                 if index != -1:
                     prompt = prompt[:index] + partial_thought
                 else:
-                    # Fallback
-                    prompt += content
+                    # Fallback if search fails: just append the thought
+                    prompt += partial_thought
             else:
                 # All other cases
                 prompt += content