Update app.py

app.py

@@ -6,6 +6,7 @@ import markdown
 from markdown.extensions.codehilite import CodeHiliteExtension
 import markdown.extensions.fenced_code
 import mistune
+import bleach
 
 API_URL = "https://host.palple.polrambora.com/pmsq"
 
@@ -123,60 +124,51 @@ def render_message(history):
     <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()
+    seen_messages = set()  # Track (role, message) pairs to avoid duplicates
+
+    allowed_tags = [
+        'p', 'strong', 'em', 'ul', 'ol', 'li', 'a', 'code', 'pre', 'br', 'blockquote', 'hr', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'img'
+    ]
+    allowed_attributes = {
+        '*': ['class', 'style'],
+        'a': ['href', 'title'],
+        'img': ['src', 'alt', 'title', 'width', 'height']
+    }
+
     for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
-        if user_message and user_message not in seen_messages:
-            seen_messages.add(user_message)
-            user_message_html = escape_and_format(user_message)
+        if user_message and ("user", user_message) not in seen_messages:
+            seen_messages.add(("user", user_message))
+            user_message_html = markdown.markdown(
+                escape_html(user_message),
+                extensions=["fenced_code", "codehilite"]
+            )
+            user_message_html = bleach.clean(user_message_html, tags=allowed_tags, attributes=allowed_attributes, strip=True)
             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; white-space: pre-wrap;'>{user_message_html}</span>
+            <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;'>
+                <div style='color: white; white-space: pre-wrap;'>{user_message_html}</div>
             </div>"""
 
-        if assistant_message and assistant_message not in seen_messages:
-            seen_messages.add(assistant_message)
-            assistant_message_html = escape_and_format(assistant_message)
+        if assistant_message and ("assistant", assistant_message) not in seen_messages:
+            seen_messages.add(("assistant", assistant_message))
+            assistant_message_html = markdown.markdown(
+                escape_html(assistant_message),
+                extensions=["fenced_code", "codehilite"]
+            )
+            assistant_message_html = bleach.clean(assistant_message_html, tags=allowed_tags, attributes=allowed_attributes, strip=True)
             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; white-space: pre-wrap;'>{assistant_message_html}</span>
+            <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;'>
+                <div style='color: white; white-space: pre-wrap;'>{assistant_message_html}</div>
             </div>"""
 
     messages_html += "</div></div>"
     return messages_html
 
 
-def escape_and_format(text):
-    allowed_tags = {'<b>', '</b>', '<i>', '</i>', '<h1>', '</h1>', '<h2>', '</h2>', '<h3>', '</h3>', 
-                    '<code>', '</code>', '<br>', '<ul>', '</ul>', '<li>', '</li>', '<blockquote>', '</blockquote>', '<p>', '</p>'}
-    
-    markdown = mistune.create_markdown(renderer="html")
-    
-    html = markdown(text)
-    
-    filtered_html = ""
-    i = 0
-    while i < len(html):
-        if html[i] == '<':
-            tag_end = html.find('>', i) + 1
-            if tag_end == 0: 
-                filtered_html += f"&#{ord(html[i])};"
-                i += 1
-                continue
-            
-            tag = html[i:tag_end]
-            if tag in allowed_tags:
-                filtered_html += tag
-            else:
-                filtered_html += ''.join(f"&#{ord(c)};" for c in tag)
-            i = tag_end
-        else:
-            filtered_html += f"&#{ord(html[i])};"
-            i += 1
-    
-    return filtered_html
-
+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;}
 """