Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import logging
-from datetime import datetime
+import json
+import os
 
 # Initialize the InferenceClient with the model ID from Hugging Face
 client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta")
@@ -13,18 +14,27 @@ logging.basicConfig(
     format='%(asctime)s - %(levelname)s - %(message)s',
 )
 
+API_KEYS_FILE = 'api_keys.json'
+
+def load_api_keys():
+    """Load API keys from the storage."""
+    if os.path.exists(API_KEYS_FILE):
+        with open(API_KEYS_FILE, 'r') as f:
+            return json.load(f)
+    return {}
+
+def authenticate(api_key: str):
+    """Authenticates the API key by checking against stored keys."""
+    api_keys = load_api_keys()
+    return api_key in api_keys.values()
+
 def log_conversation(user_message, bot_response):
-    """
-    Logs the conversation between the user and the AI.
-    
-    Args:
-        user_message (str): The user's input message.
-        bot_response (str): The AI's response.
-    """
+    """Logs the conversation between the user and the AI."""
     logging.info(f"User: {user_message}")
     logging.info(f"Bot: {bot_response}")
 
 def respond(
+    api_key: str,
     message: str,
     history: list[tuple[str, str]],
     system_message: str,
@@ -34,22 +44,12 @@ def respond(
     stop_sequence: str,
     stream_response: bool,
 ):
-    """
-    Generates a response from the AI model based on the user's message and chat history.
-    
-    Args:
-        message (str): The user's input message.
-        history (list): A list of tuples representing the conversation history (user, assistant).
-        system_message (str): A system-level message guiding the AI's behavior.
-        max_tokens (int): The maximum number of tokens for the output.
-        temperature (float): Sampling temperature for controlling the randomness.
-        top_p (float): Top-p (nucleus sampling) for controlling diversity.
-        stop_sequence (str): A custom stop sequence to end the response generation.
-        stream_response (bool): Whether to stream the response or return it as a whole.
+    """Generates a response from the AI model based on the user's message and chat history."""
 
-    Yields:
-        str: The AI's response as it is generated.
-    """
+    # Authenticate the API key
+    if not authenticate(api_key):
+        yield "Invalid API key. Access denied."
+        return
     
     # Prepare the conversation history for the API call
     messages = [{"role": "system", "content": system_message}]
@@ -95,7 +95,7 @@ def respond(
             yield response
 
     except Exception as e:
-        error_message = f"An error occurred: {str(e)}"
+        error_message = f"An error occurred: {str(e)})"
         logging.error(error_message)
         yield error_message
 
@@ -103,6 +103,7 @@ def respond(
 demo = gr.ChatInterface(
     fn=respond,
     additional_inputs=[
+        gr.Textbox(value="", label="API Key", lines=1, type="password"),
         gr.Textbox(value="You are a friendly Chatbot.", label="System Message", lines=2),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens"),
         gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),