Update app.py

app.py

@@ -1,83 +1,106 @@
 import os
 import time
 import random
+import json
 import httpx
 from typing import Dict
 
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import StreamingResponse, Response
+from fastapi.responses import Response, StreamingResponse, JSONResponse
 
 app = FastAPI()
 
-# ── CONFIG ─────────────────────────────────────
+# ─────────────────────────────────────────────
+# CONFIG
+# ─────────────────────────────────────────────
 MASTER_API_KEY = os.getenv("MASTER_API_KEY", "changeme")
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
 
-# ── LOAD KEYS ──────────────────────────────────
+# ─────────────────────────────────────────────
+# LOAD KEYS (OLLAMA_KEY_1, OLLAMA_KEY_2, ...)
+# ─────────────────────────────────────────────
 def load_keys(prefix="OLLAMA_KEY_"):
     keys = []
     for k, v in os.environ.items():
         if k.startswith(prefix) and v:
             keys.append(v.strip())
-    return sorted(keys)
+    return keys
 
 API_KEYS = load_keys()
+print(f"[INIT] Loaded {len(API_KEYS)} keys")
 
-# ── STATE ──────────────────────────────────────
+# ─────────────────────────────────────────────
+# KEY STATE
+# ─────────────────────────────────────────────
 key_status: Dict[str, Dict] = {
-    k: {"fail": 0, "cooldown": 0, "status": "active"}
+    k: {
+        "fail": 0,
+        "cooldown": 0,
+        "status": "active"
+    }
     for k in API_KEYS
 }
 
-# ── AUTH ───────────────────────────────────────
+# ─────────────────────────────────────────────
+# AUTH
+# ─────────────────────────────────────────────
 def auth(req: Request):
     if req.headers.get("Authorization") != f"Bearer {MASTER_API_KEY}":
         raise HTTPException(401, "Unauthorized")
 
-# ── KEY PICK ───────────────────────────────────
+# ─────────────────────────────────────────────
+# PICK KEY
+# ─────────────────────────────────────────────
 def pick_key():
     now = time.time()
+
     valid = [
         k for k, v in key_status.items()
         if v["status"] != "dead" and v["cooldown"] < now
     ]
+
     return random.choice(valid) if valid else None
 
-# ── STATUS ─────────────────────────────────────
+# ─────────────────────────────────────────────
+# STATUS HANDLER
+# ─────────────────────────────────────────────
 def mark_limit(k):
     key_status[k]["cooldown"] = time.time() + 60
+    key_status[k]["status"] = "limited"
 
 def mark_dead(k):
     key_status[k]["status"] = "dead"
 
 def mark_ok(k):
     key_status[k]["fail"] = 0
+    key_status[k]["status"] = "active"
 
-# ── UNIVERSAL PROXY ────────────────────────────
+# ─────────────────────────────────────────────
+# UNIVERSAL OPENAI PROXY (/v1/*)
+# ─────────────────────────────────────────────
 @app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(req: Request, path: str):
     auth(req)
 
-    # build target url dinamis
     target_url = f"{BASE_URL}/v1/{path}"
 
-    # ambil body (kalau ada)
     body = await req.body()
 
-    # copy headers kecuali auth (kita ganti)
+    # clean headers
     headers = dict(req.headers)
     headers.pop("host", None)
+    headers.pop("content-length", None)
 
     for _ in range(len(API_KEYS)):
         key = pick_key()
+
         if not key:
-            return Response("All keys exhausted", status_code=503)
+            return JSONResponse({"error": "all keys exhausted"}, 503)
 
         headers["Authorization"] = f"Bearer {key}"
 
         try:
             async with httpx.AsyncClient(timeout=None) as client:
-
                 r = await client.request(
                     method=req.method,
                     url=target_url,
@@ -86,44 +109,64 @@ async def proxy(req: Request, path: str):
                     params=req.query_params
                 )
 
-                # SUCCESS
-                if r.status_code < 400:
-                    mark_ok(key)
-
-                    # STREAMING
-                    if "text/event-stream" in r.headers.get("content-type", ""):
-                        return StreamingResponse(
-                            r.aiter_raw(),
-                            media_type="text/event-stream"
-                        )
-
-                    return Response(
-                        content=r.content,
-                        status_code=r.status_code,
-                        headers=dict(r.headers)
-                    )
+            # ─────────────────────────────
+            # SUCCESS
+            # ─────────────────────────────
+            if r.status_code < 400:
+                mark_ok(key)
 
-                # RATE LIMIT
-                elif r.status_code == 429:
-                    mark_limit(key)
+                content_type = r.headers.get("content-type", "")
 
-                # SERVER ERROR
-                elif r.status_code >= 500:
-                    key_status[key]["fail"] += 1
-                    if key_status[key]["fail"] >= 3:
-                        mark_dead(key)
+                # streaming (SSE)
+                if "text/event-stream" in content_type:
+                    return StreamingResponse(
+                        r.aiter_raw(),
+                        media_type="text/event-stream"
+                    )
+
+                return Response(
+                    content=r.content,
+                    status_code=r.status_code,
+                    headers=dict(r.headers)
+                )
 
-        except Exception:
+            # ─────────────────────────────
+            # RATE LIMIT
+            # ─────────────────────────────
+            if r.status_code == 429:
+                mark_limit(key)
+
+            # ─────────────────────────────
+            # SERVER ERROR
+            # ─────────────────────────────
+            if r.status_code >= 500:
+                key_status[key]["fail"] += 1
+                if key_status[key]["fail"] >= 3:
+                    mark_dead(key)
+
+        except Exception as e:
             key_status[key]["fail"] += 1
+
             if key_status[key]["fail"] >= 3:
                 mark_dead(key)
 
-    return Response("All keys failed", status_code=500)
+            print(f"[ERROR] key={key[:6]} err={str(e)}")
+
+    return JSONResponse(
+        {
+            "error": "all keys failed",
+            "hint": "check API keys or BASE_URL"
+        },
+        500
+    )
 
-# ── ROOT ───────────────────────────────────────
+# ─────────────────────────────────────────────
+# HEALTH CHECK
+# ─────────────────────────────────────────────
 @app.get("/")
-def home():
+def root():
     return {
-        "status": "universal proxy running",
-        "keys": len(API_KEYS)
+        "status": "ok",
+        "keys": len(API_KEYS),
+        "mode": "openai-compatible-proxy"
     }