Update app.py

app.py

@@ -10,41 +10,57 @@ sessions = {}
 PRIMARY_SYSTEM_INSTRUCTIONS = "You are P-MSQ (Messaging Service Query), a friendly AI Chatbot that can help in any situations"
 ASSISTANT_PIC_PATH = "https://huggingface.co/spaces/PLRMB/P-MSQ-API-PREVIEW/resolve/main/API.png"  
 USER_PIC_PATH = "https://huggingface.co/spaces/PLRMB/P-MSQ-API-PREVIEW/resolve/main/usr.png"
-
 def render_avatars(userid):
     try:
         response = requests.post(
             'https://host.palple.polrambora.com/userexistence',
-            json={'userid': userid},
+            json={
+                'userid': userid
+            },
             timeout=10
         )
+
         if response.status_code == 200:
             response_json = response.json()
             return response_json["avatar"]["link"]
-    except Exception:
+    except Exception as e:
         return None
+    except requests.exceptions.Timeout:
+        return None
+
 
 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)
+        response = requests.post(
+            "https://host.palple.polrambora.com/check_key_impv",
+            json=test_data,
+        )
+
         if response.status_code == 200:
+            response_json = response.json()
             avatar = render_avatars(user) or USER_PIC_PATH
+            
             if api_key not in sessions:
                 sessions[api_key] = {
-                    "history": [],
-                    "headers": {
+                    "history": [], 
+                    "headers": {   
                         "authorization": api_key,
                         "Content-Type": 'application/json'
                     },
                     "avatar": avatar,
-                    "system_message": system_message
+                    "system_message": system_message 
                 }
             return True
         elif response.status_code == 403:
             return 403
-        return False
-    except Exception:
+        else:
+            return False
+    except Exception as e:
         return False
 
 def respond(message, api_key, max_tokens, top_p, temperature):
@@ -52,12 +68,23 @@ def respond(message, api_key, max_tokens, top_p, temperature):
     history = session.get("history", [])
     headers = session.get("headers", {})
     system_message = session.get("system_message", PRIMARY_SYSTEM_INSTRUCTIONS)
-    messages = [
-        {"role": "user", "content": user_message, "profile": user_profile, "picture": user_pic}
-        if user_message else
-        {"role": "assistant", "content": assistant_message, "profile": assistant_profile, "picture": assistant_pic}
-        for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history
-    ]
+    messages = []
+    for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
+        if user_message:
+            messages.append({
+                "role": "user",
+                "content": user_message,
+                "profile": user_profile,
+                "picture": user_pic
+            })
+        if assistant_message:
+            messages.append({
+                "role": "assistant",
+                "content": assistant_message,
+                "profile": assistant_profile,
+                "picture": assistant_pic
+            })
+
     data = {
         "preferences": {
             "max_char": max_tokens,
@@ -68,46 +95,79 @@ def respond(message, api_key, max_tokens, top_p, temperature):
         "conversation_history": messages,
         "input": message
     }
+    
     response = requests.post(API_URL, headers=headers, data=json.dumps(data))
+    
     if response.status_code == 200:
         response_json = response.json()
         assistant_reply = response_json["msq"]["message"][0]
         history.append((message, assistant_reply, "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH))
-        sessions[api_key]["history"] = history
-        return history, assistant_reply
-    return history, "Error: " + response.json().get("error", "Unknown error occurred.")
+        sessions[api_key]["history"] = history  
+        return history, assistant_reply  
+    else:
+        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(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 style='color: white;'>{user_message_html}</span>
+                <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 style='color: white;'>{assistant_message_html}</span>
+                <span id="message-content" style='color: white;'>{assistant_message_html}</span>
             </div><br>"""
-    messages_html += "</div></div>"
+
+    messages_html += """
+        </div></div>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/dompurify/2.3.3/purify.min.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);
+        </script>
+        """ 
     return messages_html
 
-math_css = """
-<script type="text/javascript" src="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.js"></script>
-<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/dist/katex.min.css">
+
+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)")
+ }
 """
 
-with gr.Blocks(css=math_css) as demo:
-    with gr.Row():
+with gr.Blocks(css=".chatbox {height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex;  flex-direction: column-reverse;}", js=js) as demo:
+
+    with gr.Column(visible=True) as auth_view:
+        gr.Markdown("## P-MSQ Authorization")
+        gr.Markdown("P-MSQ is in closed alpha test! The model, api and more are subject to change.")
+        api_user_input = gr.Textbox(placeholder="snowflake", label="UserID", type='email')
+        api_key_input = gr.Textbox(placeholder="Enter your API key", label="Token", type='password')
+        auth_button = gr.Button("Authorize")
+        auth_status = gr.Textbox(label="Authorization Status", interactive=False)
+
+    with gr.Column(visible=False) as chat_view:
         gr.Markdown("## P-MSQ Chat Interface")
-        
+
         chatbot_output = gr.HTML(elem_id="chatbox-container")
 
         msg_input = gr.Text(
@@ -215,4 +275,4 @@ with gr.Blocks(css=math_css) as demo:
     save_instructions_btn.click(save_custom_instructions, inputs=[api_key_input, system_instructions_input], outputs=auth_status)
     demo.launch(show_api=False)
 if __name__ == "__main__":
-    demo.queue = False
+    demo.queue = False