.

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import requests
 import os
 import json
+import re
+import html
 
 # Set up the API endpoint and key
 API_URL = os.getenv("RUNPOD_API_URL")
@@ -12,6 +14,29 @@ headers = {
     "Content-Type": "application/json"
 }
 
+def escape_html(text):
+    return html.escape(text)
+
+def format_response_for_display(text):
+    # Escape HTML entities
+    text = escape_html(text)
+    
+    # Format <thinking> and <reflection> tags (case-insensitive)
+    text = re.sub(r'(?i)(<thinking>)(.*?)(</thinking>)', 
+                  r'<span style="font-family: monospace; color: blue;">\1</span><i>\2</i><span style="font-family: monospace; color: blue;">\3</span>', 
+                  text, flags=re.DOTALL)
+    text = re.sub(r'(?i)(<reflection>)(.*?)(</reflection>)', 
+                  r'<span style="font-family: monospace; color: green;">\1</span><i>\2</i><span style="font-family: monospace; color: green;">\3</span>', 
+                  text, flags=re.DOTALL)
+    
+    # Remove <output> tags but keep content
+    text = re.sub(r'(?i)<output>(.*?)</output>', r'\1', text, flags=re.DOTALL)
+    
+    # Replace newlines with <br> tags
+    text = text.replace('\n', '<br>')
+    
+    return text
+
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     if system_message is not None:
         messages = [{"role": "system", "content": system_message}]
@@ -47,7 +72,10 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         print(f"Formatted API Response: {json.dumps(response_json, indent=2)}")
         
         if 'choices' in response_json and len(response_json['choices']) > 0:
-            return response_json['choices'][0]['message']['content']
+            content = response_json['choices'][0]['message']['content']
+            formatted_content = format_response_for_display(content)
+            print(f"Formatted content for display: {formatted_content}")  # For debugging
+            return formatted_content
         else:
             return f"Error: Unexpected response format. Full response: {response_json}"
     
@@ -72,6 +100,7 @@ demo = gr.ChatInterface(
         gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
+    css=".message-wrap { white-space: pre-wrap; }"
 )
 
 if __name__ == "__main__":