Update app.py

app.py

@@ -10,15 +10,15 @@ from fastapi.responses import Response, StreamingResponse, JSONResponse
 
 app = FastAPI()
 
-# ─────────────────────────────────────────────
+# ─────────────────────────────
 # CONFIG
-# ─────────────────────────────────────────────
+# ─────────────────────────────
 MASTER_API_KEY = os.getenv("MASTER_API_KEY", "changeme")
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
 
-# ─────────────────────────────────────────────
-# LOAD KEYS (OLLAMA_KEY_1, OLLAMA_KEY_2, ...)
-# ─────────────────────────────────────────────
+# ─────────────────────────────
+# LOAD KEYS
+# ─────────────────────────────
 def load_keys(prefix="OLLAMA_KEY_"):
     keys = []
     for k, v in os.environ.items():
@@ -29,41 +29,35 @@ def load_keys(prefix="OLLAMA_KEY_"):
 API_KEYS = load_keys()
 print(f"[INIT] Loaded {len(API_KEYS)} keys")
 
-# ─────────────────────────────────────────────
-# KEY STATE
-# ─────────────────────────────────────────────
+# ─────────────────────────────
+# 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
-# ─────────────────────────────────────────────
+# ─────────────────────────────
 def auth(req: Request):
     if req.headers.get("Authorization") != f"Bearer {MASTER_API_KEY}":
         raise HTTPException(401, "Unauthorized")
 
-# ─────────────────────────────────────────────
-# PICK KEY
-# ─────────────────────────────────────────────
+# ─────────────────────────────
+# KEY PICKER
+# ─────────────────────────────
 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 HANDLER
-# ─────────────────────────────────────────────
+# ─────────────────────────────
 def mark_limit(k):
     key_status[k]["cooldown"] = time.time() + 60
     key_status[k]["status"] = "limited"
@@ -75,25 +69,22 @@ def mark_ok(k):
     key_status[k]["fail"] = 0
     key_status[k]["status"] = "active"
 
-# ─────────────────────────────────────────────
-# UNIVERSAL OPENAI PROXY (/v1/*)
-# ─────────────────────────────────────────────
+# ─────────────────────────────
+# UNIVERSAL OPENAI PROXY
+# ─────────────────────────────
 @app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(req: Request, path: str):
     auth(req)
 
     target_url = f"{BASE_URL}/v1/{path}"
-
     body = await req.body()
 
-    # 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 JSONResponse({"error": "all keys exhausted"}, 503)
 
@@ -101,72 +92,61 @@ async def proxy(req: Request, path: str):
 
         try:
             async with httpx.AsyncClient(timeout=None) as client:
-                r = await client.request(
+
+                # 🔥 STREAM SAFE MODE (FIX UTAMA)
+                async with client.stream(
                     method=req.method,
                     url=target_url,
                     headers=headers,
                     content=body,
                     params=req.query_params
-                )
+                ) as r:
+
+                    # ERROR RESPONSE
+                    if r.status_code >= 400:
+                        err = await r.aread()
+                        if r.status_code == 429:
+                            mark_limit(key)
+                        elif r.status_code >= 500:
+                            key_status[key]["fail"] += 1
+                            if key_status[key]["fail"] >= 3:
+                                mark_dead(key)
 
-            # ─────────────────────────────
-            # SUCCESS
-            # ─────────────────────────────
-            if r.status_code < 400:
-                mark_ok(key)
+                        continue
 
-                content_type = r.headers.get("content-type", "")
+                    mark_ok(key)
 
-                # streaming (SSE)
-                if "text/event-stream" in content_type:
+                    # STREAMING RESPONSE
                     return StreamingResponse(
-                        r.aiter_raw(),
-                        media_type="text/event-stream"
+                        r.aiter_bytes(),
+                        media_type=r.headers.get(
+                            "content-type",
+                            "application/json"
+                        )
                     )
 
-                return Response(
-                    content=r.content,
-                    status_code=r.status_code,
-                    headers=dict(r.headers)
-                )
-
-            # ─────────────────────────────
-            # 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)
 
-            print(f"[ERROR] key={key[:6]} err={str(e)}")
+            print(f"[ERROR] key={key[:6]} err={e}")
 
     return JSONResponse(
         {
             "error": "all keys failed",
-            "hint": "check API keys or BASE_URL"
+            "hint": "check keys or model access"
         },
         500
     )
 
-# ─────────────────────────────────────────────
+# ─────────────────────────────
 # HEALTH CHECK
-# ─────────────────────────────────────────────
+# ─────────────────────────────
 @app.get("/")
 def root():
     return {
         "status": "ok",
-        "keys": len(API_KEYS),
-        "mode": "openai-compatible-proxy"
+        "mode": "openai-grade-proxy",
+        "keys": len(API_KEYS)
     }