Update app.py

app.py

@@ -13,33 +13,49 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 
-# Helper function to format the streaming tokens with custom HTML/Markdown styling
+# Helper function to parse both <think> and <answer> tag structures
 def format_reasoning(text):
     """
-    Parses <think> tags and applies custom styles.
-    Uses HTML <details> for a clean, distinct 'accordion' look for reasoning.
+    Parses <think>...</think> and <answer>...</answer> tags 
+    to apply highly distinct, clean visual wrappers.
     """
+    # Extract Thinking Content
+    thinking_html = ""
     if "<think>" in text:
         if "</think>" in text:
-            # Thinking phase is complete
-            parts = text.split("</think>")
-            thinking_content = parts[0].replace("<think>", "").strip()
-            answer_content = parts[1].strip()
-            
-            return f"""<details open><summary style="color: #ff7a00; font-weight: bold; cursor: pointer;">🤔 Thinking Process (Click to collapse)</summary>
-<div style="color: #666; font-style: italic; background-color: #f9f9f9; padding: 10px; border-left: 3px solid #ff7a00; margin: 5px 0 15px 0;">
+            thinking_content = text.split("<think>")[1].split("</think>")[0].strip()
+            thinking_html = f"""<details open><summary style="color: #ff7a00; font-weight: bold; cursor: pointer; margin-bottom: 4px;">🤔 Thinking Process (Click to collapse)</summary>
+<div style="color: #555555; font-style: italic; background-color: #fdfdfd; padding: 12px; border-left: 4px solid #ff7a00; border-radius: 4px; margin: 4px 0 16px 0; font-size: 0.95em; line-height: 1.5;">
 {thinking_content}
-</div></details>
-
-{answer_content}"""
+</div></details>"""
         else:
-            # Currently thinking
-            thinking_content = text.replace("<think>", "").strip()
-            return f"""<details open><summary style="color: #ff7a00; font-weight: bold;">🤔 Thinking...</summary>
-<div style="color: #666; font-style: italic; background-color: #f9f9f9; padding: 10px; border-left: 3px solid #ff7a00; margin: 5px 0 15px 0;">
+            # Currently streaming thought process
+            thinking_content = text.split("<think>")[1].strip()
+            thinking_html = f"""<details open><summary style="color: #ff7a00; font-weight: bold;">🤔 Thinking...</summary>
+<div style="color: #555555; font-style: italic; background-color: #fdfdfd; padding: 12px; border-left: 4px solid #ff7a00; border-radius: 4px; margin: 4px 0 16px 0; font-size: 0.95em; line-height: 1.5;">
 {thinking_content}
 </div></details>"""
-    
+
+    # Extract Answer Content
+    answer_html = ""
+    if "<answer>" in text:
+        if "</answer>" in text:
+            answer_content = text.split("<answer>")[1].split("</answer>")[0].strip()
+            answer_html = f"""<div style="background-color: #f4f6f8; border-left: 4px solid #2e4053; padding: 16px; border-radius: 4px; color: #1c2833; box-shadow: inset 0 1px 3px rgba(0,0,0,0.05); line-height: 1.6;">
+{answer_content}
+</div>"""
+        else:
+            # Currently streaming final answer output
+            answer_content = text.split("<answer>")[1].strip()
+            answer_html = f"""<div style="background-color: #f4f6f8; border-left: 4px solid #2e4053; padding: 16px; border-radius: 4px; color: #1c2833; line-height: 1.6;">
+<span style="color: #2e4053; font-weight: bold; display: block; margin-bottom: 8px;">📝 Writing Response...</span>
+{answer_content}
+</div>"""
+            
+    # Combine them for the streaming UI wrapper display
+    if thinking_html or answer_html:
+        return f"{thinking_html}\n{answer_html}"
+        
     return text
 
 # 2. Main Generation Logic Handler
@@ -54,13 +70,19 @@ def generate_response(message, history, system_prompt, temperature, top_p, top_k
         if role == "user":
             prompt += f"User: {content}\n"
         elif role == "assistant":
-            # Strip out our custom HTML formatting back to raw tokens for context history
+            # Strip out our custom HTML formatting back to raw tags for context history
             clean_content = content
-            if "🤔" in content:
-                # Reconstruct rough raw tags for the model's history context
-                clean_content = re.sub(r'<details.*?>.*?<\/summary>', '<think>', clean_content)
-                clean_content = clean_content.replace('</div></details>', '</think>')
-                clean_content = re.sub(r'<div.*?>', '', clean_content)
+            if "🤔" in content or "background-color" in content:
+                # RegEx conversion logic to revert beautiful layout code back to raw tags for model's context awareness
+                import re
+                # Check for completed or uncompleted think elements
+                think_match = re.search(r'<div style="color: #555555;.*?>(.*?)</div>', clean_content, re.DOTALL)
+                ans_match = re.search(r'</div></details>\s*<div style="background-color: #f4f6f8;.*?>.*?</span>(.*?)</div>', clean_content, re.DOTALL) or re.search(r'</div></details>\s*<div style="background-color: #f4f6f8;.*?>(.*?)</div>', clean_content, re.DOTALL)
+                
+                raw_think = f"<think>{think_match.group(1).strip()}</think>" if think_match else ""
+                raw_ans = f"<answer>{ans_match.group(1).strip()}</answer>" if ans_match else ""
+                clean_content = f"{raw_think}\n{raw_ans}".strip()
+                
             prompt += f"Assistant: {clean_content}\n"
     
     # Inject active raw message with Quasar's default target sequence patterns
@@ -94,7 +116,6 @@ with gr.Blocks() as demo:
     with gr.Row():
         # Left Workspace Column: Chat Layout
         with gr.Column(scale=3):
-            # REMOVED type="messages" to solve the init crash
             chatbot = gr.Chatbot(height=500)
             msg_input = gr.Textbox(placeholder="Enter reasoning prompt...", label="Your Message")
             with gr.Row():