Update app.py

app.py

@@ -1,11 +1,23 @@
+import bleach
 import gradio as gr
 import requests
 import json
 import os
 import markdown
 from markdown.extensions.codehilite import CodeHiliteExtension
+from markdown.extensions.fenced_code import FencedCodeExtension
 import markdown.extensions.fenced_code
-
+import html
+import re
+
+ALLOWED_TAGS = [
+    'b', 'i', 'strong', 'em', 'a', 'p', 'ul', 'ol', 'li', 'br', 'hr',
+    'blockquote', 'code', 'pre', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'img'
+]
+ALLOWED_ATTRIBUTES = {
+    'a': ['href', 'title', 'target'],
+    'img': ['src', 'alt', 'title', 'width', 'height']
+}
 API_URL = "https://host.palple.polrambora.com/pmsq"
 
 sessions = {}
@@ -32,15 +44,16 @@ def render_avatars(userid):
 
 
 def authorize(user, api_key, system_message):
-    test_data = {"user": user, "key": api_key}
+    test_data = {
+        "user": user,
+        "key": api_key
+    }
+
     try:
         response = requests.post(
             "https://host.palple.polrambora.com/check_key_impv",
             json=test_data,
         )
-        print(f"Request Data: {test_data}")
-        print(f"Response Status: {response.status_code}")
-        print(f"Response Text: {response.text}")
 
         if response.status_code == 200:
             response_json = response.json()
@@ -62,10 +75,8 @@ def authorize(user, api_key, system_message):
         else:
             return False
     except Exception as e:
-        print(f"Exception occurred: {str(e)}")
         return False
 
-
 def respond(message, api_key, max_tokens, top_p, temperature):
     session = sessions.get(api_key, {})
     headers = session.get("headers", {})
@@ -97,7 +108,8 @@ def respond(message, api_key, max_tokens, top_p, temperature):
         response = requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True)
         
         if response.status_code == 200:
-            for line in response.iter_lines(decode_unicode=True):
+            for line in response.iter_lines():
+                line = line.decode('utf-8')
                 if line.strip():
                     try:
                         if line.startswith("data:"):
@@ -116,19 +128,30 @@ def respond(message, api_key, max_tokens, top_p, temperature):
     except Exception as e:
         yield f"Error: {str(e)}"
 
+def format_code_blocks(text):
+   
+    code_pattern = re.compile(r'```(.*?)```', re.DOTALL)
+    return code_pattern.sub(lambda match: f"<pre><code>{html.escape(match.group(1))}</code></pre>", text)
+
+def render_markdown(text):
+    return markdown.markdown(text, extensions=[FencedCodeExtension()])
+
+
 
+import re
 
 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;">"""
 
-    seen_messages = set()  # Track (role, message) pairs to avoid duplicates
+    seen_messages = set()
 
     for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
         if user_message and ("user", user_message) not in seen_messages:
             seen_messages.add(("user", user_message))
-            user_message_html = markdown.markdown(user_message, extensions=["fenced_code", "codehilite"])
+            user_message_html = process_message(user_message)
             messages_html += f"""
             <div style='display: flex; flex-direction: column; align-items: flex-start; margin-bottom: 10px;'>
                 <img src='{user_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-bottom: 5px;'>
@@ -137,7 +160,7 @@ def render_message(history):
 
         if assistant_message and ("assistant", assistant_message) not in seen_messages:
             seen_messages.add(("assistant", assistant_message))
-            assistant_message_html = markdown.markdown(assistant_message, extensions=["fenced_code", "codehilite"])
+            assistant_message_html = process_message(assistant_message)
             messages_html += f"""
             <div style='display: flex; flex-direction: column; align-items: flex-start; margin-bottom: 10px;'>
                 <img src='{assistant_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-bottom: 5px;'>
@@ -147,11 +170,36 @@ def render_message(history):
     messages_html += "</div></div>"
     return messages_html
 
+def process_message(message):
+    code_block_pattern = re.compile(r"```(\w+)?\n(.*?)```", re.DOTALL)
+    parts = code_block_pattern.split(message)
+
+    for i in range(len(parts)):
+        if i % 3 == 0:  
+            parts[i] = escape_html(parts[i]) 
+            parts[i] = render_inline_markdown(parts[i])  
+        elif i % 3 == 1: 
+            continue
+        else:  
+            lang_class = parts[i - 1] or ""  
+            parts[i] = f"<pre><code class='language-{lang_class}'>{escape_html(parts[i])}</code></pre>"
+
+    return "".join(parts)
+
+
+def escape_html(text):
+    return (text
+        .replace("&", "&amp;")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+        .replace('"', "&quot;")
+        .replace("'", "&#x27;"))
+
+def render_inline_markdown(text):
+    text = markdown.markdown(text, output_format='xhtml')
+    return text
 
 
-def escape_html(unsafe_text):
-    escaped_text = ''.join(f"&#{ord(char)};" if char not in ('\n', '\r') else '<br>' for char in unsafe_text)
-    return escaped_text
 
 css="""
 .chatbox {height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex;  flex-direction: column-reverse;}
@@ -169,10 +217,19 @@ with gr.Blocks(css=css) as demo:
 
     with gr.Column(visible=False) as chat_view:
         gr.HTML("""
-    <script type="text/javascript" id="MathJax-script" async
-        src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/3.2.2/es5/tex-mml-chtml.js">
-    </script>
-    """)
+<link href="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/themes/prism.css" rel="stylesheet">
+<script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/prism.min.js"></script>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/components/prism-python.min.js"></script>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/components/prism-javascript.min.js"></script>
+<script>
+    MathJax = {
+        tex: { inlineMath: [['\\(', '\\)'], ['\\[', '\\]']] },
+        svg: { scale: 1.2 }
+    };
+</script>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/3.2.2/es5/tex-svg.min.js"></script>
+""")
+
         gr.Markdown("## P-MSQ Chat Interface")
         chatbot_output = gr.HTML(elem_id="chatbox-container")
         gr.Markdown(elem_id="chatbox-container") 
@@ -214,6 +271,7 @@ with gr.Blocks(css=css) as demo:
 
 
 
+
         
         def regenerate_response(history, last_message, max_tokens, top_p, temperature):
             return "", []