Update app.py

app.py

@@ -10,7 +10,7 @@ API_URL = "https://host.palple.polrambora.com/pmsq"
 
 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"  
+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):
@@ -25,8 +25,6 @@ def render_avatars(userid):
             return response_json["avatar"]["link"]
     except Exception:
         return None
-    except requests.exceptions.Timeout:
-        return None
 
 def authorize(user, api_key, system_message):
     test_data = {"user": user, "key": api_key}
@@ -36,17 +34,16 @@ def authorize(user, api_key, system_message):
             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:
@@ -89,18 +86,15 @@ def respond(message, api_key, max_tokens, top_p, temperature):
         "input": message
     }
 
-    response = requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True)
-    if response.status_code == 200:
-        assistant_reply = ""
-        for chunk in response.iter_content(chunk_size=1024):
-            if chunk:
-                assistant_reply += chunk.decode()
-                yield assistant_reply
-        history.append((message, assistant_reply, "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH))
-        sessions[api_key]["history"] = history  
-        yield history, assistant_reply
-    else:
-        yield history, "Error: " + response.json().get("error", "Unknown error occurred.")
+    with requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True) as response:
+        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
+        else:
+            return history, "Error: " + response.json().get("error", "Unknown error occurred.")
 
 def render_message(history):
     messages_html = """
@@ -128,6 +122,12 @@ def render_message(history):
 
     return messages_html
 
+def user_interaction(message, history, api_key, max_tokens, top_p, temperature):
+    loading_message = history + [(message, "Loading...", "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH)]
+    yield render_message(loading_message), loading_message, ""
+    history, assistant_reply = respond(message, api_key, max_tokens, top_p, temperature)
+    yield render_message(history), history, ""
+
 def escape_html(unsafe_text):
     return (
         unsafe_text.replace("&", "&")
@@ -144,56 +144,49 @@ css = """
 with gr.Blocks(css=css) 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.")
+        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.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>
-    """)
         gr.Markdown("## P-MSQ Chat Interface")
         chatbot_output = gr.HTML(elem_id="chatbox-container")
-        gr.Markdown(elem_id="chatbox-container") 
-        msg_input = gr.Text(
-            show_label=False, 
-            placeholder="Type your message and press Shift+Enter...", 
-            lines=2,
-            elem_id="input-text"
-        )
+        msg_input = gr.Text(show_label=False, placeholder="Type your message and press Shift+Enter...", lines=2)
         send_btn = gr.Button("Send")
         regen_btn = gr.Button("Clear")
-        system_instructions_input = gr.Textbox(placeholder="Enter custom instructions (optional)", 
-                                               label="Custom System Instructions", 
-                                               lines=2)
-        save_instructions_btn = gr.Button("Save Instructions")        
-        gr.Markdown("### Settings")
+        system_instructions_input = gr.Textbox(placeholder="Enter custom instructions (optional)", label="Custom System Instructions", lines=2)
+        save_instructions_btn = gr.Button("Save Instructions")
         max_tokens = gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens")
         top_p = gr.Slider(minimum=0, maximum=2, value=0.8, step=0.1, label="Top P")
         temperature = gr.Slider(minimum=0.1, maximum=1, value=0.7, step=0.1, label="Temperature")
-
-        history_state = gr.State([])  
-        last_message_state = gr.State("")  
+        history_state = gr.State([])
+        last_message_state = gr.State("")
 
         msg_input.submit(
-            respond, 
-            inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature], 
+            user_interaction,
+            inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature],
             outputs=[chatbot_output, history_state, msg_input],
         )
         send_btn.click(
-            respond, 
-            inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature], 
+            user_interaction,
+            inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature],
             outputs=[chatbot_output, history_state, msg_input],
         )
-        regen_btn.click(lambda x: ("", []), inputs=[api_key_input], outputs=[chatbot_output, history_state])
-
-    auth_button.click(authorize, inputs=[api_user_input, api_key_input], outputs=[auth_view, chat_view])
+        regen_btn.click(lambda _: ("", []), outputs=[chatbot_output, history_state])
 
-demo.launch(show_api=False)
+    auth_button.click(
+        lambda user, key: (authorize(user, key, PRIMARY_SYSTEM_INSTRUCTIONS), ""),
+        inputs=[api_user_input, api_key_input],
+        outputs=[chat_view, auth_status],
+    )
+    save_instructions_btn.click(
+        lambda key, instructions: ("", ""),
+        inputs=[api_key_input, system_instructions_input],
+        outputs=[auth_status],
+    )
+    demo.launch(show_api=False)
 
 if __name__ == "__main__":
     demo.queue = False