Update app.py

app.py

@@ -10,47 +10,59 @@ 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):
     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,
         )
+
         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
         else:
             return False
-    except Exception:
+    except Exception as e:
         return False
 
 def respond(message, api_key, max_tokens, top_p, temperature):
@@ -59,6 +71,7 @@ def respond(message, api_key, max_tokens, top_p, temperature):
     headers = session.get("headers", {})
     system_message = session.get("system_message", PRIMARY_SYSTEM_INSTRUCTIONS)
     messages = []
+
     for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
         if user_message:
             messages.append({
@@ -80,21 +93,38 @@ def respond(message, api_key, max_tokens, top_p, temperature):
             "max_char": max_tokens,
             "temperature": temperature,
             "top_p": top_p,
-            "system_message": system_message
+            "system_message": system_message,
+            "stream": True  
         },
         "conversation_history": messages,
         "input": message
     }
 
-    with requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True) as response:
+    try:
+        response = requests.post(API_URL, headers=headers, data=json.dumps(data), stream=True)
+
         if response.status_code == 200:
-            response_json = response.json()
-            assistant_reply = response_json["msq"]["message"][0]
+            assistant_reply = ""
+            for line in response.iter_lines(decode_unicode=True):
+                if line:
+                    try:
+                        chunk = json.loads(line)
+                        chunk_message = chunk.get("chunk", {}).get("content", "")
+                        assistant_reply += chunk_message
+
+                        yield assistant_reply
+                    except Exception as e:
+                        print(f"Stream chunk error: {e}")
+
             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.")
+            yield f"Error: {response.status_code} - {response.text}"
+    except Exception as e:
+        yield f"Error: {str(e)}"
+
+
+
 
 def render_message(history):
     messages_html = """
@@ -122,12 +152,6 @@ 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("&", "&")
@@ -137,11 +161,12 @@ def escape_html(unsafe_text):
         .replace("'", "'")
     )
 
-css = """
+css="""
 .chatbox {height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex;  flex-direction: column-reverse;}
 """
-
 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.")
@@ -151,43 +176,122 @@ with gr.Blocks(css=css) as demo:
         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")
-        msg_input = gr.Text(show_label=False, placeholder="Type your message and press Shift+Enter...", lines=2)
+        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"
+        )
+
         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")
+
+        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")
+        
         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("")  
+
+        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, ""
+
+            assistant_reply = ""
+            for partial_reply in respond(message, api_key, max_tokens, top_p, temperature):
+                partial_history = history + [(message, partial_reply, "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH)]
+                yield render_message(partial_history), partial_history, ""
+
+            history, assistant_reply = partial_history, partial_reply
+            yield render_message(history), history, ""
+            
+        def regenerate_response(history, last_message, max_tokens, top_p, temperature):
+            return "", []
+
+        def clear_history(api_key):
+            if api_key in sessions:
+                sessions[api_key]["history"] = []  
+            return "", []
+
+        def load_conversation(api_key):
+            session = sessions.get(api_key, {})
+            history = session.get("history", [])
+            return render_message(history), history
 
         msg_input.submit(
-            user_interaction,
-            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(
-            user_interaction,
-            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 _: ("", []), outputs=[chatbot_output, history_state])
 
-    auth_button.click(
-        lambda user, key: ("Authorized" if authorize(user, key, PRIMARY_SYSTEM_INSTRUCTIONS) else "Authorization Failed"),
-        inputs=[api_user_input, api_key_input],
-        outputs=[auth_status],
-    )
+        regen_btn.click(clear_history, 
+                        inputs=[api_key_input], 
+                        outputs=[chatbot_output, history_state])
 
-    save_instructions_btn.click(
-        lambda key, instructions: ("", ""),
-        inputs=[api_key_input, system_instructions_input],
-        outputs=[auth_status],
-    )
-    demo.launch(show_api=False)
+    with gr.Column(visible=False) as blacklist_view:
+        gr.Markdown("## P-MSQ Authorization")
+        gr.Markdown("Your linked ID appears to be blacklisted, and your API Key is pending on removal, if you believe this is a mistake, please try reaching us out.")
+    
+    
+    def authorize_and_proceed(user, api_key):
+        if authorize(user, api_key, PRIMARY_SYSTEM_INSTRUCTIONS):
+            gr.Info("Loading, please wait.")
+            messages_html, history = load_conversation(api_key)
 
+            return (
+            gr.update(visible=False), 
+            gr.update(visible=True),  
+            messages_html,
+            history
+            )
+        elif authorize(user, api_key, PRIMARY_SYSTEM_INSTRUCTIONS) == 403: 
+            return (
+                gr.update(visible=False),  
+                gr.update(visible=False), 
+                gr.update(visible=True), 
+            )
+        else:
+            gr.Warning("Incorrect userid/token")
+            return (
+                gr.update(visible=True),  
+                gr.update(visible=False), 
+                auth_status.update(value="Invalid userid/token")
+            )
+
+    def save_custom_instructions(api_key, custom_instructions):
+        if api_key in sessions:
+            gr.Info("Instructions updated, we recommend to start the new conversation to make it more efficient.")
+            sessions[api_key]["system_message"] = custom_instructions  
+            return "Instructions updated!", gr.update(value="")
+        else:
+            gr.Warning("Your session has been expired, please refresh the page and login again.")
+            return "Session not found.", gr.update(value="")
+
+
+    
+
+    auth_button.click(authorize_and_proceed, inputs=[api_user_input, api_key_input], outputs=[auth_view, chat_view, chatbot_output, history_state])
+    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