Update main.py

main.py

@@ -18,8 +18,21 @@ parser.add_argument("--port", type=int, help="Set the port.(default: 7860)", def
 args = parser.parse_args()
 
 base_url = os.getenv('MODEL_BASE_URL')
+API_KEY = os.getenv("API_KEY")  # Load API key from environment variable
 print(base_url)
 
+def verify_api_key():
+    """Check if the API key is valid."""
+    api_key = request.headers.get("Authorization")
+    if not api_key or api_key != f"Bearer {API_KEY}":
+        return jsonify({"error": "Unauthorized"}), 401
+
+@app.before_request
+def before_request():
+    """Apply API key validation to all routes except index."""
+    if request.endpoint not in ['index']:
+        return verify_api_key()
+
 @app.route('/api/v1/models', methods=["GET", "POST"])
 @app.route('/v1/models', methods=["GET", "POST"])
 def model_list():
@@ -27,18 +40,8 @@ def model_list():
     model_list = {
         "object": "list",
         "data": [
-            {
-                "id": "hunyuan-turbo",
-                "object": "model",
-                "created": time_now,
-                "owned_by": "tastypear"
-            },
-            {
-                "id": "gpt-3.5-turbo",
-                "object": "model",
-                "created": time_now,
-                "owned_by": "tastypear"
-            }
+            {"id": "hunyuan-turbo", "object": "model", "created": time_now, "owned_by": "tastypear"},
+            {"id": "gpt-3.5-turbo", "object": "model", "created": time_now, "owned_by": "tastypear"}
         ]
     }
     return jsonify(model_list)
@@ -56,16 +59,9 @@ def index():
 @app.route("/v1/chat/completions", methods=["POST", "OPTIONS"])
 def chat_completions():
     if request.method == "OPTIONS":
-        return Response(
-            headers={
-                "Access-Control-Allow-Origin": "*",
-                "Access-Control-Allow-Headers": "*",
-            }
-        )
-
+        return Response(headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Headers": "*"})
+    
     data = request.get_json()
-
-    # Set a default system prompt.
     system = "You are a helpful assistant who responds in English."
     chat_history = []
     prompt = ""
@@ -73,8 +69,6 @@ def chat_completions():
     if "messages" in data:
         messages = data["messages"]
         message_size = len(messages)
-
-        # Get the last message's content as the main prompt.
         prompt = messages[-1].get("content")
         
         for i in range(message_size - 1):
@@ -83,45 +77,32 @@ def chat_completions():
             if role_this == "system":
                 system = messages[i].get("content")
             elif role_this == "user":
-                if role_next == "assistant":
-                    chat_history.append([messages[i].get("content"), messages[i + 1].get("content")])
-                else:
-                    chat_history.append([messages[i].get("content"), " "])
+                chat_history.append([messages[i].get("content"), messages[i + 1].get("content") if role_next == "assistant" else " "])
         
-        # Combine system prompt with the user's prompt.
         if system:
             prompt = system + "\n" + prompt
         
         fn_index = 3
-
-        # Generate a random 10-character hash.
-        chars = string.ascii_lowercase + string.digits
-        session_hash = "".join(random.choice(chars) for _ in range(10))
+        session_hash = "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(10))
         
         single_prompt_data = {
-            'data': [
-                prompt,
-                [],
-            ],
+            'data': [prompt, []],
             'event_data': None,
             'fn_index': 1,
             'trigger_id': 5,
             'session_hash': session_hash,
         }
-        response = requests.post(f'{base_url}/gradio_api/run/predict', json=single_prompt_data)
+        requests.post(f'{base_url}/gradio_api/run/predict', json=single_prompt_data)
         
         context_data = {
-            'data': [
-                None,
-                chat_history + [[prompt, None]]
-            ],
+            'data': [None, chat_history + [[prompt, None]]],
             'event_data': None,
             'fn_index': fn_index,
             'trigger_id': 5,
             'session_hash': session_hash,
         }
-        response = requests.post(f"{base_url}/gradio_api/queue/join", json=context_data)
-
+        requests.post(f"{base_url}/gradio_api/queue/join", json=context_data)
+    
     def generate():
         url = f"{base_url}/gradio_api/queue/data?session_hash={session_hash}"
         data = requests.get(url, stream=True)
@@ -132,22 +113,13 @@ def chat_completions():
                 decoded_line = line.decode("utf-8")
                 json_line = json.loads(decoded_line[6:])
                 if json_line["msg"] == "process_starts":
-                    res_data = gen_res_data({}, time_now=time_now, start=True)
-                    yield f"data: {json.dumps(res_data)}\n\n"
+                    yield f"data: {json.dumps(gen_res_data({}, time_now, start=True))}\n\n"
                 elif json_line["msg"] == "process_generating":
-                    res_data = gen_res_data(json_line, time_now=time_now)
-                    yield f"data: {json.dumps(res_data)}\n\n"
+                    yield f"data: {json.dumps(gen_res_data(json_line, time_now))}\n\n"
                 elif json_line["msg"] == "process_completed":
                     yield "data: [DONE]"
 
-    return Response(
-        generate(),
-        mimetype="text/event-stream",
-        headers={
-            "Access-Control-Allow-Origin": "*",
-            "Access-Control-Allow-Headers": "*",
-        },
-    )
+    return Response(generate(), mimetype="text/event-stream", headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Headers": "*"})
 
 def gen_res_data(data, time_now=0, start=False):
     res_data = {
@@ -160,8 +132,8 @@ def gen_res_data(data, time_now=0, start=False):
     if start:
         res_data["choices"][0]["delta"] = {"role": "assistant", "content": ""}
     else:
-        chat_pair = data["output"]["data"][0]
-        if chat_pair == []:
+        chat_pair = data.get("output", {}).get("data", [])[0]
+        if not chat_pair:
             res_data["choices"][0]["finish_reason"] = "stop"
         else:
             res_data["choices"][0]["delta"] = {"content": chat_pair[-1][-1]}