Update main.py

main.py

@@ -1,3 +1,9 @@
+#!/usr/bin/env python3
+# ─── PATCH SSL EARLY ─────────────────────────────────────────────────────────────
+from gevent import monkey
+monkey.patch_all()
+
+# ─── STANDARD IMPORTS ───────────────────────────────────────────────────────────
 import os
 import random
 import string
@@ -6,10 +12,8 @@ import json
 import requests
 from functools import wraps
 from flask import Flask, request, Response, jsonify
-import gevent.pywsgi
-from gevent import monkey; monkey.patch_all()
 
-# —— Load & check env vars ——
+# ─── ENV & CONFIG ────────────────────────────────────────────────────────────────
 API_KEY = os.getenv("API_KEY")
 if not API_KEY:
     raise RuntimeError("Missing API_KEY env var")
@@ -20,17 +24,16 @@ if not MODEL_BASE_URL:
 
 SPACE_URL = os.getenv("SPACE_URL", "")
 
-# —— Flask setup ——
 app = Flask(__name__)
 app.json.sort_keys = False
 
-# —— Error handler to show real errors ——
+# ─── GLOBAL ERROR HANDLER ────────────────────────────────────────────────────────
 @app.errorhandler(Exception)
 def handle_all_errors(e):
     app.logger.exception(e)
     return jsonify({"error": str(e)}), 500
 
-# —— API‑key decorator ——
+# ─── API‑KEY DECORATOR ───────────────────────────────────────────────────────────
 def require_api_key(f):
     @wraps(f)
     def decorated(*args, **kwargs):
@@ -44,7 +47,7 @@ def require_api_key(f):
         return f(*args, **kwargs)
     return decorated
 
-# —— Model list endpoint ——
+# ─── MODEL LIST ─────────────────────────────────────────────────────────────────
 @app.route('/api/v1/models', methods=['GET', 'POST'])
 @app.route('/v1/models',      methods=['GET', 'POST'])
 @require_api_key
@@ -58,7 +61,7 @@ def model_list():
         ]
     })
 
-# —— Home page ——
+# ─── INDEX ──────────────────────────────────────────────────────────────────────
 @app.route("/", methods=["GET"])
 def index():
     return Response(
@@ -67,7 +70,7 @@ def index():
         f'Full API: {SPACE_URL}/api/v1/chat/completions'
     )
 
-# —— Chat completions ——
+# ─── CHAT COMPLETIONS ────────────────────────────────────────────────────────────
 @app.route("/api/v1/chat/completions", methods=["POST", "OPTIONS"])
 @app.route("/v1/chat/completions",      methods=["POST", "OPTIONS"])
 @require_api_key
@@ -78,43 +81,39 @@ def chat_completions():
         )
 
     data = request.get_json() or {}
-    messages = data.get("messages")
-    if not messages:
+    msgs = data.get("messages")
+    if not msgs:
         return jsonify({"error": "Missing 'messages' field"}), 400
 
-    # parse messages
     system = "You are a helpful assistant."
-    chat_history = []
-    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","")])
+    history = []
+    for i, m in enumerate(msgs[:-1]):
+        if m.get("role") == "system":
+            system = m.get("content", system)
+        elif m.get("role") == "user":
+            nxt = msgs[i+1].get("role")
+            if nxt == "assistant":
+                history.append([m.get("content",""), msgs[i+1].get("content","")])
             else:
-                chat_history.append([msg.get("content",""), ""])
+                history.append([m.get("content",""), ""])
 
-    prompt = messages[-1].get("content","")
+    prompt = msgs[-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}
+    payload = {"data":[prompt, history, system], "fn_index":0, "session_hash":session_hash}
 
     def generate():
-        # enqueue
-        requests.post(f"{MODEL_BASE_URL}/queue/join", json=json_prompt)
+        requests.post(f"{MODEL_BASE_URL}/queue/join", json=payload)
         url = f"{MODEL_BASE_URL}/queue/data?session_hash={session_hash}"
         resp = requests.get(url, stream=True)
-        start_time = int(time.time())
+        start_ts = int(time.time())
 
         for line in resp.iter_lines():
-            if not line: 
-                continue
+            if not line: continue
             msg = json.loads(line.decode("utf-8")[6:])
             if msg["msg"] == "process_starts":
-                chunk = gen_res_data({}, start=True, time_now=start_time)
+                chunk = make_chunk({}, start=True, ts=start_ts)
             elif msg["msg"] == "process_generating":
-                chunk = gen_res_data(msg, start=False, time_now=start_time)
+                chunk = make_chunk(msg, start=False, ts=start_ts)
             elif msg["msg"] == "process_completed":
                 yield "data: [DONE]"
                 break
@@ -128,26 +127,27 @@ def chat_completions():
         headers={"Access-Control-Allow-Origin":"*","Access-Control-Allow-Headers":"*"},
     )
 
-def gen_res_data(data, start=False, time_now=None):
-    if time_now is None:
-        time_now = int(time.time())
+def make_chunk(data, start=False, ts=None):
+    if ts is None:
+        ts = int(time.time())
     base = {
         "id": "chatcmpl",
         "object": "chat.completion.chunk",
-        "created": time_now,
+        "created": ts,
         "model": "glm-4",
         "choices": [{"index": 0, "finish_reason": None}]
     }
     if start:
         base["choices"][0]["delta"] = {"role": "assistant", "content": ""}
     else:
-        chat_pair = data.get("output",{}).get("data",[None,None])[1]
-        if not chat_pair:
+        pair = data.get("output",{}).get("data",[None,None])[1] or []
+        if not pair:
             base["choices"][0]["finish_reason"] = "stop"
         else:
-            base["choices"][0]["delta"] = {"content": chat_pair[-1][-1]}
+            base["choices"][0]["delta"] = {"content": pair[-1][-1]}
     return base
 
+# ─── RUN ────────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser()
@@ -155,7 +155,7 @@ if __name__ == "__main__":
     parser.add_argument("--port", type=int, default=7860)
     args = parser.parse_args()
 
-    # turn on debug so you see errors in browser
+    # Dev: use Flask’s debug server so you see tracebacks in browser
     app.run(host=args.host, port=args.port, debug=True)
-    # once it works, you can switch back to:
+    # Prod: swap to gevent server
     # gevent.pywsgi.WSGIServer((args.host, args.port), app).serve_forever()