Update app.py

app.py

@@ -115,42 +115,63 @@ def render_message(history):
 
     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)
+            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 id="message-content" style='color: white;'>{user_message_html}</span>
+                <span style='color: white;'>{user_message_html}</span>
             </div><br>"""
 
         if assistant_message:
-            assistant_message_html = markdown.markdown(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 id="message-content" style='color: white;'>{assistant_message_html}</span>
+                <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://cdnjs.cloudflare.com/ajax/libs/mathjax/3.2.2/es5/tex-mml-chtml.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;");
+                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
+            // 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>
-        """ 
+    """
     return messages_html
 
 
+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;")
+    )
+
 js = """
 function HTMLClean() {
     alert("This page is in-dev stage, expect lots of restarts, the interface is open-sourced (not the backend that is responsible for API Keys)")
@@ -174,37 +195,6 @@ 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(