Update app.py

app.py

@@ -12,41 +12,32 @@ 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 as e:
+    except Exception:
         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": [], 
@@ -62,7 +53,7 @@ def authorize(user, api_key, system_message):
             return 403
         else:
             return False
-    except Exception as e:
+    except Exception:
         return False
 
 def respond(message, api_key, max_tokens, top_p, temperature):
@@ -97,18 +88,19 @@ 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))
-    
+
+    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 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  
-        return history, assistant_reply  
+        yield history, assistant_reply
     else:
-        return history, "Error: " + response.json().get("error", "Unknown error occurred.")
-
+        yield history, "Error: " + response.json().get("error", "Unknown error occurred.")
 
 def render_message(history):
     messages_html = """
@@ -145,15 +137,14 @@ 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.")
+        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")
@@ -174,16 +165,13 @@ with gr.Blocks(css=css) as demo:
             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")        
         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")
@@ -191,86 +179,21 @@ with gr.Blocks(css=css) as demo:
         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, ""
-
-            history, assistant_reply = respond(message, api_key, max_tokens, top_p, temperature)
-            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, 
+            respond, 
             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, 
+            respond, 
             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])
 
-        regen_btn.click(clear_history, 
-                        inputs=[api_key_input], 
-                        outputs=[chatbot_output, history_state])
-
-    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, inputs=[api_user_input, api_key_input], outputs=[auth_view, chat_view])
 
-    
+demo.launch(show_api=False)
 
-    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