Update main.py

main.py

@@ -1,57 +1,74 @@
-import gevent.pywsgi
-from gevent import monkey;monkey.patch_all()
-from flask import Flask, request, Response, jsonify
-import argparse
-import requests
+import os
 import random
 import string
 import time
 import json
-import os
+import requests
+from functools import wraps
+from flask import Flask, request, Response, jsonify, abort
+import gevent.pywsgi
+from gevent import monkey; monkey.patch_all()
 
 app = Flask(__name__)
 app.json.sort_keys = False
 
-parser = argparse.ArgumentParser(description="An example of Zhipu GLM-4 with a similar API to OAI.")
-parser.add_argument("--host", type=str, help="Set the ip address.(default: 0.0.0.0)", default='0.0.0.0')
-parser.add_argument("--port", type=int, help="Set the port.(default: 7860)", default=7860)
+# ——— Load and enforce API key ———
+API_KEY = os.getenv("API_KEY")
+if not API_KEY:
+    raise RuntimeError("Missing API_KEY in env")
+
+def require_api_key(f):
+    @wraps(f)
+    def decorated(*args, **kwargs):
+        # look in X-API-Key header or Authorization: Bearer ...
+        key = request.headers.get("X-API-Key") or request.headers.get("Authorization", "")
+        if key.startswith("Bearer "):
+            key = key.split(" ", 1)[1]
+        if not key:
+            return jsonify({"error": "API key missing"}), 401
+        if key != API_KEY:
+            return jsonify({"error": "Invalid API key"}), 403
+        return f(*args, **kwargs)
+    return decorated
+
+# ——— Your existing arg parsing & base_url ———
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--host", default="0.0.0.0")
+parser.add_argument("--port", type=int, default=7860)
 args = parser.parse_args()
 
-base_url = os.getenv('MODEL_BASE_URL')
+base_url = os.getenv("MODEL_BASE_URL")
+
 
 @app.route('/api/v1/models', methods=["GET", "POST"])
-@app.route('/v1/models', methods=["GET", "POST"])
+@app.route('/v1/models',      methods=["GET", "POST"])
+@require_api_key
 def model_list():
     time_now = int(time.time())
-    model_list = {
+    models = {
         "object": "list",
         "data": [
-            {
-                "id": "glm-4",
-                "object": "model",
-                "created": time_now,
-                "owned_by": "tastypear"
-            },
-            {
-                "id": "gpt-3.5-turbo",
-                "object": "model",
-                "created": time_now,
-                "owned_by": "tastypear"
-            }
+            {"id": "glm-4",       "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)
+    return jsonify(models)
+
 
 @app.route("/", methods=["GET"])
 def index():
-    return Response(f'ZhipuAI GLM-4 OpenAI Compatible API<br><br>'+
-        f'Set "{os.getenv("SPACE_URL")}/api" as proxy (or API Domain) in your Chatbot.<br><br>'+
-        f'The complete API is: {os.getenv("SPACE_URL")}/api/v1/chat/completions')
+    return Response(
+        f'ZhipuAI GLM-4 OpenAI Compatible API<br><br>'
+        f'Set "{os.getenv("SPACE_URL")}/api" as proxy in your Chatbot.<br><br>'
+        f'Full API: {os.getenv("SPACE_URL")}/api/v1/chat/completions'
+    )
+
 
 @app.route("/api/v1/chat/completions", methods=["POST", "OPTIONS"])
-@app.route("/v1/chat/completions", methods=["POST", "OPTIONS"])
+@app.route("/v1/chat/completions",      methods=["POST", "OPTIONS"])
+@require_api_key
 def chat_completions():
-
     if request.method == "OPTIONS":
         return Response(
             headers={
@@ -60,66 +77,56 @@ def chat_completions():
             }
         )
 
-    data = request.get_json()
+    data = request.get_json() or {}
+    if "messages" not in data:
+        return jsonify({"error": "Missing 'messages' field"}), 400
 
     # reorganize data
     system = "You are a helpful assistant."
     chat_history = []
-    prompt = ""
-    
-    if "messages" in data:
-        messages = data["messages"]
-        message_size = len(messages)
-
-        prompt = messages[-1].get("content")
-        for i in range(message_size - 1):
-            role_this = messages[i].get("role")
-            role_next = messages[i + 1].get("role")
-            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"), " "])
-
-        # print(f'{system = }')
-        # print(f'{chat_history = }')
-        # print(f'{prompt = }')
-
-        fn_index = 0
-
-        # gen a random char(11) hash
-        chars = string.ascii_lowercase + string.digits
-        session_hash = "".join(random.choice(chars) for _ in range(11))
-
-        json_prompt = {
-            "data": [prompt, chat_history, system],
-            "fn_index": fn_index,
-            "session_hash": session_hash,
-        }
+    messages = data["messages"]
+    prompt = messages[-1].get("content", "")
+
+    for i in range(len(messages) - 1):
+        r0 = messages[i].get("role")
+        r1 = messages[i+1].get("role")
+        if r0 == "system":
+            system = messages[i]["content"]
+        elif r0 == "user":
+            if r1 == "assistant":
+                chat_history.append([messages[i]["content"], messages[i+1]["content"]])
+            else:
+                chat_history.append([messages[i]["content"], " "])
+
+    # random session id
+    session_hash = "".join(random.choices(string.ascii_lowercase + string.digits, k=11))
+    json_prompt = {
+        "data": [prompt, chat_history, system],
+        "fn_index": 0,
+        "session_hash": session_hash,
+    }
 
     def generate():
-        response = requests.post(f"{base_url}/queue/join", json=json_prompt)
+        # enqueue job
+        requests.post(f"{base_url}/queue/join", json=json_prompt)
         url = f"{base_url}/queue/data?session_hash={session_hash}"
-        data = requests.get(url, stream=True)
-
-        time_now = int(time.time())
-
-        for line in data.iter_lines():
-            if line:
-                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"
-                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"
-                elif json_line["msg"] == "process_completed":
-                    yield "data: [DONE]"
+        stream = requests.get(url, stream=True)
+
+        start_time = int(time.time())
+        for line in stream.iter_lines():
+            if not line: 
+                continue
+            msg = json.loads(line.decode("utf-8")[6:])
+            if msg["msg"] == "process_starts":
+                chunk = gen_res_data({}, time_now=start_time, start=True)
+            elif msg["msg"] == "process_generating":
+                chunk = gen_res_data(msg, time_now=start_time)
+            elif msg["msg"] == "process_completed":
+                yield "data: [DONE]"
+                break
+            else:
+                continue
+            yield f"data: {json.dumps(chunk)}\n\n"
 
     return Response(
         generate(),
@@ -132,25 +139,23 @@ def chat_completions():
 
 
 def gen_res_data(data, time_now=0, start=False):
-    res_data = {
+    base = {
         "id": "chatcmpl",
         "object": "chat.completion.chunk",
         "created": time_now,
         "model": "glm-4",
-        "choices": [{"index": 0, "finish_reason": None}],
+        "choices": [{"index": 0, "finish_reason": None}]
     }
-
     if start:
-        res_data["choices"][0]["delta"] = {"role": "assistant", "content": ""}
+        base["choices"][0]["delta"] = {"role": "assistant", "content": ""}
     else:
         chat_pair = data["output"]["data"][1]
-        if chat_pair == []:
-            res_data["choices"][0]["finish_reason"] = "stop"
+        if not chat_pair:
+            base["choices"][0]["finish_reason"] = "stop"
         else:
-            res_data["choices"][0]["delta"] = {"content": chat_pair[-1][-1]}
-    return res_data
+            base["choices"][0]["delta"] = {"content": chat_pair[-1][-1]}
+    return base
 
 
 if __name__ == "__main__":
-    # app.run(host=args.host, port=args.port, debug=True)
     gevent.pywsgi.WSGIServer((args.host, args.port), app).serve_forever()