Update app.py

app.py

@@ -79,15 +79,17 @@ if not API_KEYS:
 # =====================================================
 # KEY STATUS & ROUND ROBIN
 # =====================================================
+RATE_LIMIT_COOLDOWN = int(os.getenv("RATE_LIMIT_COOLDOWN", "62"))  # detik cooldown setelah rate limit
+
 key_status = {}
 for idx, k in enumerate(API_KEYS, 1):
     key_status[k] = {
         "index": idx,
         "prefix": k[:8] + "..." if len(k) > 8 else k,
-        "healthy": True,
         "busy": False,
         "success": 0,
         "fail": 0,
+        "rate_limited_until": 0.0,   # timestamp epoch; 0 = tidak sedang cooldown
     }
 
 rr_index = 0
@@ -222,50 +224,66 @@ def is_rate_limited_error_body(text: str) -> bool:
 # =====================================================
 # KEY MANAGEMENT
 # =====================================================
-async def get_key(exclude=None):
+def _get_available_key(exclude: set) -> str | None:
+    """
+    Internal (sync, dipanggil dalam _key_lock): cari key yang:
+    1. Tidak sedang busy
+    2. Tidak sedang cooldown rate limit
+    3. Tidak ada di exclude set
+    Round-robin.
+    """
     global rr_index
+    now = time.time()
+    for _ in range(len(API_KEYS)):
+        rr_index = (rr_index + 1) % len(API_KEYS)
+        key = API_KEYS[rr_index]
+        st = key_status[key]
+        if not st["busy"] and now >= st["rate_limited_until"] and key not in exclude:
+            st["busy"] = True
+            return key
+    return None
+
+def _next_available_time() -> float:
+    """Kapan key pertama keluar dari cooldown (epoch seconds)."""
+    now = time.time()
+    times = [st["rate_limited_until"] for st in key_status.values() if st["rate_limited_until"] > now]
+    return min(times) if times else now
+
+async def get_key(exclude=None):
     if exclude is None:
         exclude = set()
-
     async with _key_lock:
-        # Reset semua key kalau semuanya unhealthy
-        if not any(v["healthy"] for v in key_status.values()):
-            log("⚠️ All API Keys unhealthy. Resetting all...")
-            for v in key_status.values():
-                v["fail"] = 0
-                v["healthy"] = True
-
-        for _ in range(len(API_KEYS)):
-            rr_index = (rr_index + 1) % len(API_KEYS)
-            key = API_KEYS[rr_index]
-            st = key_status[key]
-
-            if st["healthy"] and not st["busy"] and key not in exclude:
-                st["busy"] = True
-                return key
-
-    return None
+        return _get_available_key(exclude)
 
 async def release_key(key):
     async with _key_lock:
         if key in key_status:
             key_status[key]["busy"] = False
 
+async def mark_rate_limited(key):
+    """Tandai key kena rate limit: set cooldown RATE_LIMIT_COOLDOWN detik."""
+    async with _key_lock:
+        if key in key_status:
+            until = time.time() + RATE_LIMIT_COOLDOWN
+            key_status[key]["rate_limited_until"] = until
+            key_status[key]["fail"] += 1
+            idx = key_status[key]["index"]
+            log(f"⏳ key#{idx} cooldown {RATE_LIMIT_COOLDOWN}s (sampai {time.strftime('%H:%M:%S', time.localtime(until))})")
+
 async def mark_fail(key):
     async with _key_lock:
         if key in key_status:
             key_status[key]["fail"] += 1
-            if key_status[key]["fail"] >= 3:
-                key_status[key]["healthy"] = False
 
 async def mark_ok(key):
     async with _key_lock:
         if key in key_status:
             key_status[key]["success"] += 1
             key_status[key]["fail"] = 0
-            key_status[key]["healthy"] = True
+            key_status[key]["rate_limited_until"] = 0.0
 
-async def wait_for_free_key(exclude=None, max_wait=60.0, interval=0.3):
+async def wait_for_free_key(exclude=None, max_wait=120.0, interval=0.5):
+    """Tunggu key tersedia, max max_wait detik."""
     elapsed = 0.0
     while elapsed < max_wait:
         key = await get_key(exclude)
@@ -277,21 +295,39 @@ async def wait_for_free_key(exclude=None, max_wait=60.0, interval=0.3):
 
 async def get_key_infinite(exclude=None):
     """
-    Tunggu key tanpa batas waktu (infinite). Reset exclude setiap full cycle.
-    Dipakai supaya tidak pernah berhenti mencoba.
+    Tunggu key tanpa batas waktu (infinite).
+    - Kalau ada key tersedia: return langsung.
+    - Kalau semua key busy/cooldown: sleep TEPAT sampai key paling cepat ready,
+      lalu retry — tidak perlu hammering setiap 2 detik.
+    - exclude di-reset setiap full cycle supaya key bisa dipakai lagi.
     """
     local_exclude = set(exclude) if exclude else set()
     cycle = 0
+
     while True:
-        key = await get_key(exclude=local_exclude)
-        if key:
-            return key, local_exclude
+        async with _key_lock:
+            key = _get_available_key(local_exclude)
+            if key:
+                return key, local_exclude
+
+            # Hitung berapa lama sampai key berikutnya ready
+            now = time.time()
+            next_ready = _next_available_time()
+            wait_sec = max(0.5, next_ready - now)
+
+            all_in_cooldown = all(
+                st["rate_limited_until"] > now or st["busy"]
+                for st in key_status.values()
+            )
 
-        # Semua key dicoba, mulai cycle baru
-        await asyncio.sleep(1.0)  # jeda 1 detik antar cycle
-        cycle += 1
-        log(f"🔄 Key cycle #{cycle}, all keys tried/busy, retrying...")
-        local_exclude.clear()  # reset exclude supaya semua key dicoba lagi
+        if all_in_cooldown:
+            cycle += 1
+            log(f"⏳ Semua key cooldown. Tunggu {wait_sec:.1f}s sampai key berikutnya ready... (cycle #{cycle})")
+            local_exclude.clear()  # reset exclude agar key dicoba lagi setelah cooldown
+            await asyncio.sleep(wait_sec)
+        else:
+            # Ada key yang sudah lewat cooldown tapi mungkin busy — tunggu sebentar
+            await asyncio.sleep(0.3)
 
 
 # =====================================================
@@ -469,11 +505,14 @@ def openai_response_to_anthropic(data: dict, original_model: str) -> dict:
 @app.get("/")
 async def root():
     async with _key_lock:
+        now = time.time()
         keys_info = {}
         for k, v in key_status.items():
+            rl_until = v["rate_limited_until"]
+            cooldown_remaining = max(0, rl_until - now)
             keys_info[v["prefix"]] = {
-                "status": "BUSY" if v["busy"] else "IDLE",
-                "healthy": v["healthy"],
+                "status": "BUSY" if v["busy"] else ("COOLDOWN" if cooldown_remaining > 0 else "IDLE"),
+                "cooldown_remaining_sec": round(cooldown_remaining, 1) if cooldown_remaining > 0 else 0,
                 "success": v["success"],
                 "fail": v["fail"],
             }
@@ -484,6 +523,7 @@ async def root():
         "base_url": CEREBRAS_BASE_URL,
         "default_model": DEFAULT_MODEL,
         "max_request_tokens": MAX_REQUEST_TOKENS,
+        "rate_limit_cooldown_sec": RATE_LIMIT_COOLDOWN,
         "total_keys": len(API_KEYS),
         "keys": keys_info,
     }
@@ -587,8 +627,7 @@ async def chat(req: Request):
 
                 if is_rate_limited_status(r.status_code) or (r.status_code != 200 and is_rate_limited_error_body(r.text)):
                     log(f"RATE LIMITED: key#{ki['index']}")
-                    await mark_fail(key)
-                    await asyncio.sleep(2)
+                    await mark_rate_limited(key)
                     continue
 
                 if r.status_code != 200:
@@ -610,6 +649,7 @@ async def chat(req: Request):
 
     # -----------------------------------------
     # STREAM — infinite loop, tidak pernah stop
+    # Ketika semua key cooldown, sleep TEPAT sampai key siap
     # -----------------------------------------
     async def stream_gen():
         exclude = set()
@@ -630,8 +670,7 @@ async def chat(req: Request):
 
                         if is_rate_limited_status(r.status_code):
                             log(f"STREAM RATE LIMITED: key#{ki['index']}")
-                            await mark_fail(key)
-                            await asyncio.sleep(2)
+                            await mark_rate_limited(key)
                             continue
 
                         if r.status_code != 200:
@@ -649,7 +688,6 @@ async def chat(req: Request):
 
                             raw = line[6:] if line.startswith("data: ") else line
 
-                            # Cek error CF hanya pada JSON error object, bukan model output
                             try:
                                 j = json.loads(raw)
                                 if "error" in j and "choices" not in j:
@@ -663,12 +701,12 @@ async def chat(req: Request):
                             yield line + "\n\n"
 
                         if hit_limit:
-                            await mark_fail(key)
+                            await mark_rate_limited(key)
                             continue
 
                         yield "data: [DONE]\n\n"
                         await mark_ok(key)
-                        return  # sukses, keluar dari loop
+                        return  # sukses
 
             except Exception as e:
                 log(f"STREAM EXCEPTION: key#{ki['index']} - {e}")
@@ -768,7 +806,7 @@ async def anthropic_messages(req: Request):
 
                 if is_rate_limited_status(r.status_code) or (r.status_code != 200 and is_rate_limited_error_body(r.text)):
                     log(f"RATE LIMITED: key#{ki['index']}")
-                    await mark_fail(key)
+                    await mark_rate_limited(key)
                     continue
 
                 if r.status_code != 200:
@@ -821,8 +859,7 @@ async def anthropic_messages(req: Request):
 
                         if is_rate_limited_status(r.status_code):
                             log(f"STREAM RATE LIMITED: key#{ki['index']}")
-                            await mark_fail(key)
-                            await asyncio.sleep(2)
+                            await mark_rate_limited(key)
                             continue
 
                         if r.status_code != 200:
@@ -953,7 +990,7 @@ async def anthropic_messages(req: Request):
                                     })
 
                         if hit_limit:
-                            await mark_fail(key)
+                            await mark_rate_limited(key)
                             continue
 
                         # Tutup text block