Update app.py

app.py

@@ -4,38 +4,65 @@ from openai import OpenAI
 from typing import List, Tuple
 
 CLIENTS = {
-    "perplexity":{"key":os.getenv('PX_KEY'),"endpoint":"https://api.perplexity.ai"},
-    "hyperbolic":{"key":os.getenv('HYPERBOLIC_XYZ_KEY'),"endpoint":"https://api.hyperbolic.xyz/v1"},
-    "huggingface":{"key":os.getenv('HF_KEY'),"endpoint":"https://huggingface.co/api/inference-proxy/together"},
+    "perplexity": {"key": os.getenv('PX_KEY'), "endpoint": "https://api.perplexity.ai"},
+    "hyperbolic": {"key": os.getenv('HYPERBOLIC_XYZ_KEY'), "endpoint": "https://api.hyperbolic.xyz/v1"},
+    "huggingface": {"key": os.getenv('HF_KEY'), "endpoint": "https://huggingface.co/api/inference-proxy/together"},
 }
 for client_type in CLIENTS:
-    CLIENTS[client_type]["client"] = OpenAI(base_url=CLIENTS[client_type]["endpoint"], api_key=CLIENTS[client_type]["key"])
+    CLIENTS[client_type]["client"] = OpenAI(
+        base_url=CLIENTS[client_type]["endpoint"],
+        api_key=CLIENTS[client_type]["key"]
+    )
 
 PASSWORD = os.getenv("PASSWD")
 
 # Define available models
 AVAILABLE_MODELS = {
-    "Llama3.3-70b-Instruct (Hyperbolic.xyz)": {"model_name":"meta-llama/Llama-3.3-70B-Instruct","type":"hyperbolic"},
-    "Llama3.1-8b-Instruct (Hyperbolic.xyz)": {"model_name":"meta-llama/Meta-Llama-3.1-8B-Instruct","type":"hyperbolic"},
-    "DeepSeek V3 (Hyperbolic.xyz)": {"model_name":"deepseek-ai/DeepSeek-V3","type":"hyperbolic"},
-    "DeepSeek V3 (HuggingFace.co)": {"model_name":"deepseek-ai/DeepSeek-V3","type":"huggingface"},
-    "Sonar Pro (Perplexity.ai)": {"model_name":"sonar-pro","type":"perplexity"},
-    "Sonar (Perplexity.ai)": {"model_name":"sonar","type":"perplexity"},
+    "Llama3.3-70b-Instruct (Hyperbolic.xyz)": {
+        "model_name": "meta-llama/Llama-3.3-70B-Instruct",
+        "type": "hyperbolic"
+    },
+    "Llama3.1-8b-Instruct (Hyperbolic.xyz)": {
+        "model_name": "meta-llama/Meta-Llama-3.1-8B-Instruct",
+        "type": "hyperbolic"
+    },
+    "DeepSeek V3 (Hyperbolic.xyz)": {
+        "model_name": "deepseek-ai/DeepSeek-V3",
+        "type": "hyperbolic"
+    },
+    "DeepSeek V3 (HuggingFace.co)": {
+        "model_name": "deepseek-ai/DeepSeek-V3",
+        "type": "huggingface"
+    },
+    "Sonar Pro (Perplexity.ai)": {
+        "model_name": "sonar-pro",
+        "type": "perplexity"
+    },
+    "Sonar (Perplexity.ai)": {
+        "model_name": "sonar",
+        "type": "perplexity"
+    },
 }
 
 def respond(
     message: str,
     history: List[Tuple[str, str]],
+    session_password: str,  # added parameter from session state
     system_message: str,
     model_choice: str,
     max_tokens: int,
     temperature: float,
     top_p: float,
 ):
-    """Handles chatbot responses with Perplexity AI."""
-    
+    """Handles chatbot responses with password re-checking."""
+    # Re-check the session password on every message
+    if session_password != PASSWORD:
+        yield "Error: Invalid session password. Please refresh the page and enter the correct password."
+        return
+
     if model_choice not in AVAILABLE_MODELS:
-        return "Error: Invalid model selection."
+        yield "Error: Invalid model selection."
+        return
 
     messages = [{"role": "system", "content": system_message}]
     for user_msg, assistant_msg in history:
@@ -49,7 +76,7 @@ def respond(
     citations = []
 
     selected_client = CLIENTS[AVAILABLE_MODELS[model_choice]["type"]]["client"]
-    
+
     try:
         stream = selected_client.chat.completions.create(
             model=AVAILABLE_MODELS[model_choice]["model_name"],
@@ -68,7 +95,7 @@ def respond(
             if hasattr(chunk, "citations") and chunk.citations:
                 citations = chunk.citations
 
-        # Append citations as clickable links
+        # Append citations as clickable links, if any
         if citations:
             citation_text = "\n\nSources:\n" + "\n".join(
                 [f"[{i+1}] [{url}]({url})" for i, url in enumerate(citations)]
@@ -79,14 +106,19 @@ def respond(
     except Exception as e:
         yield f"Error: {str(e)}"
 
+
 def check_password(input_password):
-    """Validates the password before showing the chat interface."""
+    """Validates the password and, if valid, stores it in session state."""
     if input_password == PASSWORD:
-        return gr.update(visible=False), gr.update(visible=True)
+        # Return the password to store in the session state.
+        return gr.update(visible=False), gr.update(visible=True), input_password
     else:
-        return gr.update(value="", interactive=True), gr.update(visible=False)
+        return gr.update(value="", interactive=True), gr.update(visible=False), ""
 
 with gr.Blocks() as demo:
+    # A hidden state component to store the session password
+    session_password = gr.State("")
+
     with gr.Column():
         password_input = gr.Textbox(
             type="password", label="Enter Password", interactive=True
@@ -115,15 +147,19 @@ with gr.Blocks() as demo:
             minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"
         )
 
+        # Note: The session_password is now passed as an additional input to the chat function.
         chat = gr.ChatInterface(
-            respond,
+            fn=respond,
             api_name=False,
-            chatbot=gr.Chatbot(height=400),  # Set the desired height here
-            additional_inputs=[system_prompt, model_choice, max_tokens, temperature, top_p]  # Pass extra parameters
+            chatbot=gr.Chatbot(height=400),  # Set desired height here
+            additional_inputs=[session_password, system_prompt, model_choice, max_tokens, temperature, top_p]
         )
 
+    # Now, the submit_button click updates three outputs: the password_input, chat_interface visibility, and session_password state.
     submit_button.click(
-        check_password, inputs=password_input, outputs=[password_input, chat_interface]
+        fn=check_password,
+        inputs=password_input,
+        outputs=[password_input, chat_interface, session_password]
     )
 
 if __name__ == "__main__":