Update app.py

app.py

@@ -3,12 +3,6 @@ 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"
 
@@ -114,59 +108,48 @@ def respond(message, api_key, max_tokens, top_p, temperature):
         return history, "Error: " + response.json().get("error", "Unknown error occurred.")
 
 
-def detect_content_type(message):
-    """
-    Automatically detect whether the message contains Math, HTML, or plain Markdown.
-    Returns a tuple: (is_math, is_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."""
-    return (
-        unsafe_text.replace("&", "&amp;")
-        .replace("<", "&lt;")
-        .replace(">", "&gt;")
-        .replace('"', "&quot;")
-        .replace("'", "&#039;")
-    )
+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(user_message)
+            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 id="message-content" style='color: white;'>{user_message_html}</span>
+            </div><br>"""
+
+        if assistant_message:
+            assistant_message_html = markdown.markdown(assistant_message)
+            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 id="message-content" 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://cdnjs.cloudflare.com/ajax/libs/mathjax/3.2.2/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;");
+            }
+            let message = document.getElementById('message-content').innerHTML;
+            document.getElementById('message-content').innerHTML = escapeHtml(message);
+            MathJax.typeset();  // Process math expressions
+        </script>
+        """ 
+    return messages_html
+
 
 js = """
 function HTMLClean() {
@@ -191,6 +174,37 @@ with gr.Blocks(css=css, js=js) as demo:
 
     with gr.Column(visible=False) as chat_view:
         gr.Markdown("## P-MSQ Chat Interface")
+        gr.Markdown(label="answer",
+                                    latex_delimiters=[{
+                                        "left": "\\(",
+                                        "right": "\\)",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{equation\}",
+                                        "right": "\\end\{equation\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{align\}",
+                                        "right": "\\end\{align\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{alignat\}",
+                                        "right": "\\end\{alignat\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{gather\}",
+                                        "right": "\\end\{gather\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{CD\}",
+                                        "right": "\\end\{CD\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\[",
+                                        "right": "\\]",
+                                        "display": True
+                                    }],
+                                    elem_id="chatbox-container")
         chatbot_output = gr.HTML(elem_id="chatbox-container")
 
         msg_input = gr.Text(