Update main.py

main.py

@@ -5,22 +5,35 @@ import time
 import json
 import requests
 from functools import wraps
-from flask import Flask, request, Response, jsonify, abort
+from flask import Flask, request, Response, jsonify
 import gevent.pywsgi
 from gevent import monkey; monkey.patch_all()
 
+# —— Load & check env vars ——
+API_KEY = os.getenv("API_KEY")
+if not API_KEY:
+    raise RuntimeError("Missing API_KEY env var")
+
+MODEL_BASE_URL = os.getenv("MODEL_BASE_URL")
+if not MODEL_BASE_URL:
+    raise RuntimeError("Missing MODEL_BASE_URL env var")
+
+SPACE_URL = os.getenv("SPACE_URL", "")
+
+# —— Flask setup ——
 app = Flask(__name__)
 app.json.sort_keys = False
 
-# ——— Load and enforce API key ———
-API_KEY = os.getenv("API_KEY")
-if not API_KEY:
-    raise RuntimeError("Missing API_KEY in env")
+# —— Error handler to show real errors ——
+@app.errorhandler(Exception)
+def handle_all_errors(e):
+    app.logger.exception(e)
+    return jsonify({"error": str(e)}), 500
 
+# —— API‑key decorator ——
 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]
@@ -31,96 +44,77 @@ def require_api_key(f):
         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")
-
-
-@app.route('/api/v1/models', methods=["GET", "POST"])
-@app.route('/v1/models',      methods=["GET", "POST"])
+# —— Model list endpoint ——
+@app.route('/api/v1/models', methods=['GET', 'POST'])
+@app.route('/v1/models',      methods=['GET', 'POST'])
 @require_api_key
 def model_list():
-    time_now = int(time.time())
-    models = {
+    now = int(time.time())
+    return jsonify({
         "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": now, "owned_by": "tastypear"},
+            {"id": "gpt-3.5-turbo", "object": "model", "created": now, "owned_by": "tastypear"}
         ]
-    }
-    return jsonify(models)
-
+    })
 
+# —— Home page ——
 @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 in your Chatbot.<br><br>'
-        f'Full API: {os.getenv("SPACE_URL")}/api/v1/chat/completions'
+        f'Set "{SPACE_URL}/api" as proxy in your Chatbot.<br><br>'
+        f'Full API: {SPACE_URL}/api/v1/chat/completions'
     )
 
-
+# —— Chat completions ——
 @app.route("/api/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={
-                "Access-Control-Allow-Origin": "*",
-                "Access-Control-Allow-Headers": "*",
-            }
+            headers={"Access-Control-Allow-Origin":"*","Access-Control-Allow-Headers":"*"}
         )
 
     data = request.get_json() or {}
-    if "messages" not in data:
+    messages = data.get("messages")
+    if not messages:
         return jsonify({"error": "Missing 'messages' field"}), 400
 
-    # reorganize data
+    # parse messages
     system = "You are a helpful assistant."
     chat_history = []
-    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"]])
+    for i, msg in enumerate(messages[:-1]):
+        role = msg.get("role")
+        if role == "system":
+            system = msg.get("content", system)
+        elif role == "user":
+            next_role = messages[i+1].get("role")
+            if next_role == "assistant":
+                chat_history.append([msg.get("content",""), messages[i+1].get("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,
-    }
+                chat_history.append([msg.get("content",""), ""])
 
-    def generate():
-        # enqueue job
-        requests.post(f"{base_url}/queue/join", json=json_prompt)
-        url = f"{base_url}/queue/data?session_hash={session_hash}"
-        stream = requests.get(url, stream=True)
+    prompt = messages[-1].get("content","")
+    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():
+        # enqueue
+        requests.post(f"{MODEL_BASE_URL}/queue/join", json=json_prompt)
+        url = f"{MODEL_BASE_URL}/queue/data?session_hash={session_hash}"
+        resp = requests.get(url, stream=True)
         start_time = int(time.time())
-        for line in stream.iter_lines():
+
+        for line in resp.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)
+                chunk = gen_res_data({}, start=True, time_now=start_time)
             elif msg["msg"] == "process_generating":
-                chunk = gen_res_data(msg, time_now=start_time)
+                chunk = gen_res_data(msg, start=False, time_now=start_time)
             elif msg["msg"] == "process_completed":
                 yield "data: [DONE]"
                 break
@@ -131,14 +125,12 @@ def chat_completions():
     return Response(
         generate(),
         mimetype="text/event-stream",
-        headers={
-            "Access-Control-Allow-Origin": "*",
-            "Access-Control-Allow-Headers": "*",
-        },
+        headers={"Access-Control-Allow-Origin":"*","Access-Control-Allow-Headers":"*"},
     )
 
-
-def gen_res_data(data, time_now=0, start=False):
+def gen_res_data(data, start=False, time_now=None):
+    if time_now is None:
+        time_now = int(time.time())
     base = {
         "id": "chatcmpl",
         "object": "chat.completion.chunk",
@@ -149,13 +141,21 @@ def gen_res_data(data, time_now=0, start=False):
     if start:
         base["choices"][0]["delta"] = {"role": "assistant", "content": ""}
     else:
-        chat_pair = data["output"]["data"][1]
+        chat_pair = data.get("output",{}).get("data",[None,None])[1]
         if not chat_pair:
             base["choices"][0]["finish_reason"] = "stop"
         else:
             base["choices"][0]["delta"] = {"content": chat_pair[-1][-1]}
     return base
 
-
 if __name__ == "__main__":
-    gevent.pywsgi.WSGIServer((args.host, args.port), app).serve_forever()
+    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()
+
+    # turn on debug so you see errors in browser
+    app.run(host=args.host, port=args.port, debug=True)
+    # once it works, you can switch back to:
+    # gevent.pywsgi.WSGIServer((args.host, args.port), app).serve_forever()