Update app.py

app.py

@@ -3,6 +3,12 @@ import requests
 import json
 import os
 import markdown
+import gradio as gr
+import re
+import markdown
+from pygments import highlight
+from pygments.lexers import HtmlLexer, PythonLexer, TextLexer
+from pygments.formatters import HtmlFormatter
 
 API_URL = "https://host.palple.polrambora.com/pmsq"
 
@@ -108,59 +114,49 @@ def respond(message, api_key, max_tokens, top_p, temperature):
         return history, "Error: " + response.json().get("error", "Unknown error occurred.")
 
 
-def render_message(history):
-    messages_html = """
-    <div id="chatbox-container" class="chatbox" style="height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex; flex-direction: column-reverse;">
-        <div id="messages" style="display: block; margin-bottom: 10px;">"""
-
-    for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
-        if user_message:
-            user_message_html = markdown.markdown(
-                escape_html(user_message), 
-                extensions=["fenced_code", "codehilite", "md_in_html"]
-            )
-            messages_html += f"""
-            <div style='display: flex; align-items: center; margin-bottom: 10px;'>
-                <img src='{user_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-right: 10px;'>
-                <span style='color: white;'>{user_message_html}</span>
-            </div><br>"""
-
-        if assistant_message:
-            assistant_message_html = markdown.markdown(
-                escape_html(assistant_message), 
-                extensions=["fenced_code", "codehilite", "md_in_html"]
-            )
-            messages_html += f"""
-            <div style='display: flex; align-items: center; margin-bottom: 10px;'>
-                <img src='{assistant_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-right: 10px;'>
-                <span style='color: white;'>{assistant_message_html}</span>
-            </div><br>"""
-
-    messages_html += """
-        </div></div>
-        <script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/2.3.3/purify.min.js"></script>
-        <script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
-        <script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
-        <script>
-            function escapeHtml(unsafe) {
-                return unsafe
-                    .replace(/&/g, "&amp;")
-                    .replace(/</g, "&lt;")
-                    .replace(/>/g, "&gt;")
-                    .replace(/"/g, "&quot;")
-                    .replace(/'/g, "&#039;");
-            }
-            // Sanitize messages using DOMPurify
-            let messages = document.querySelectorAll('#messages span');
-            messages.forEach((message) => {
-                message.innerHTML = DOMPurify.sanitize(message.innerHTML); // Sanitize content
-            });
-            // Render math using MathJax
-            MathJax.typeset();
-        </script>
+def detect_content_type(message):
+    """
+    Automatically detect whether the message contains Math, HTML, or plain Markdown.
+    Returns a tuple: (is_math, is_html)
     """
-    return messages_html
+    is_math = bool(re.search(r"(\$\$.*?\$\$|\$.*?\$)", message, re.DOTALL))
+    
+    is_html = bool(re.search(r"<\/?[a-z][\s\S]*?>", message))
+    
+    return is_math, is_html
 
+def render_message(message):
+    """
+    Render a message dynamically, detecting if it's Math, HTML, or Markdown with code.
+    """
+    is_math, is_html = detect_content_type(message)
+
+    if is_math:
+        return gr.HTML(f"""
+        <div style="font-size: 16px; font-family: Arial;">
+            <script type="text/javascript" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
+            {message}
+        </div>
+        """)
+    elif is_html:
+        formatter = HtmlFormatter(style="default", noclasses=True)
+        highlighted_html = highlight(message, HtmlLexer(), formatter)
+        return gr.HTML(f"""
+        <div style="background: #f5f5f5; border: 1px solid #ddd; padding: 10px; overflow-x: auto;">
+            {highlighted_html}
+        </div>
+        """)
+    else:
+        def markdown_with_code_blocks(md):
+            html = markdown.markdown(md, extensions=["fenced_code"])
+            return html
+
+        formatted_message = markdown_with_code_blocks(message)
+        return gr.HTML(f"""
+        <div style="font-size: 16px; font-family: Arial;">
+            {formatted_message}
+        </div>
+        """)
 
 def escape_html(unsafe_text):
     """Escape HTML special characters to prevent code injection."""