Update app.py

app.py

@@ -10,21 +10,22 @@ import asyncio
 app = FastAPI()
 
 # ==========================================
-# KONFIGURASI
+# KONFIGURASI & LOAD KEYS
 # ==========================================
-BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
-MASTER_API_KEY = os.getenv("MASTER_API_KEY", "ollama")
+BASE_URL = os.getenv("BASE_URL", "https://elysiadev11-proxyollma.hf.space")
+MASTER_API_KEY = os.getenv("MASTER_API_KEY", "ollama-proxy-free")
 
 OLLAMA_KEYS = []
-for i in range(1, 25):
+# Mendukung hingga 100 API Key (OLLAMA_KEY_1 sampai OLLAMA_KEY_100)
+for i in range(1, 101):
     key = os.getenv(f"OLLAMA_KEY_{i}")
     if key:
         OLLAMA_KEYS.append(key)
 
 if not OLLAMA_KEYS:
-    OLLAMA_KEYS.append("ollama")
+    OLLAMA_KEYS.append("ollama-proxy-free") # Dummy key jika ENV kosong
 
-# Status Key
+# Inisialisasi Status Key
 key_status = {}
 for idx, k in enumerate(OLLAMA_KEYS, 1):
     key_status[k] = {
@@ -39,19 +40,29 @@ for idx, k in enumerate(OLLAMA_KEYS, 1):
 def log(msg):
     print(f"[{time.strftime('%H:%M:%S')}] {msg}")
 
-def get_available_key(exclude_keys=None):
-    """Mencari key yang sehat dan sedang tidak digunakan (idle)"""
+def get_and_lock_key(exclude_keys=None):
+    """
+    Fungsi ATOMIC: Mencari key yang idle sekaligus langsung MENGUNCINYA.
+    Ini menghilangkan celah kebocoran saat burst request.
+    """
     if exclude_keys is None:
         exclude_keys = set()
         
-    available = [
-        k for k, v in key_status.items() 
-        if v["healthy"] and not v["in_use"] and k not in exclude_keys
-    ]
-    
-    if available:
-        return available[0]
-    return None
+    # Cek apakah semua key mati? Jika ya, reset semuanya agar sistem tidak macet permanen
+    if not any(v["healthy"] for v in key_status.values()):
+        log("⚠️ Semua API Key berstatus mati/unhealthy. Melakukan RESET MASSAL...")
+        for v in key_status.values():
+            v["failures"] = 0
+            v["healthy"] = True
+            
+    # Proses pencarian dan penguncian instan
+    for k, v in key_status.items():
+        if v["healthy"] and not v["in_use"] and k not in exclude_keys:
+            # ---> GEMBOK DIPASANG DI SINI (ATOMIC) <---
+            v["in_use"] = True 
+            return k
+            
+    return None # Mengembalikan None jika semua key sedang dipakai
 
 # ==========================================
 # ENDPOINTS
@@ -60,12 +71,13 @@ def get_available_key(exclude_keys=None):
 def root():
     return {
         "status": "ok",
-        "keys_loaded": len(OLLAMA_KEYS),
+        "total_keys_loaded": len(OLLAMA_KEYS),
         "keys_status": {
             v["prefix"]: {
                 "status": "BUSY" if v["in_use"] else "IDLE",
                 "healthy": v["healthy"],
-                "success": v["success"]
+                "success": v["success"],
+                "failures": v["failures"]
             } for v in key_status.values()
         }
     }
@@ -76,7 +88,7 @@ async def chat(req: Request):
     if auth_key != MASTER_API_KEY:
         return JSONResponse({"error": "Unauthorized"}, status_code=401)
     
-    # Tangkap error jika client kabur sebelum request terkirim penuh
+    # Tangkap error jika client kabur (ClientDisconnect)
     try:
         body = await req.json()
     except ClientDisconnect:
@@ -92,7 +104,7 @@ async def chat(req: Request):
     # ==========================================
     if not is_stream:
         tried_keys = set()
-        for attempt in range(len(OLLAMA_KEYS)): # Loop sebanyak jumlah key
+        for attempt in range(len(OLLAMA_KEYS)): 
             if len(tried_keys) >= len(OLLAMA_KEYS):
                 tried_keys.clear()
 
@@ -105,14 +117,16 @@ async def chat(req: Request):
                     log("Client membatalkan request saat mengantre (Non-Stream).")
                     return Response(status_code=499)
                 
-                key = get_available_key(exclude_keys=tried_keys)
-                if key: break
-                await asyncio.sleep(1)
+                # Gunakan fungsi Atomic Lock
+                key = get_and_lock_key(exclude_keys=tried_keys)
+                if key: 
+                    break # Langsung keluar loop, key SUDAH DIKUNCI
+                    
+                await asyncio.sleep(0.5) # Cek tiap setengah detik
 
             ki = key_status[key]
-            ki["in_use"] = True
             tried_keys.add(key)
-            log(f"LOCK: key#{ki['index']} (Non-Stream)")
+            log(f"LOCK ACQUIRED: key#{ki['index']} (Non-Stream)")
             
             try:
                 async with httpx.AsyncClient(timeout=120.0) as client:
@@ -123,7 +137,13 @@ async def chat(req: Request):
                     )
                 if resp.status_code == 200:
                     ki["success"] += 1
+                    ki["failures"] = 0
                     return Response(content=resp.content, media_type=resp.headers.get("content-type"))
+                elif resp.status_code == 429:
+                    ki["failures"] += 1
+                    ki["healthy"] = False
+                    log(f"RATE LIMIT: key#{ki['index']} - Skip ke key berikutnya.")
+                    continue
                 else:
                     ki["failures"] += 1
                     continue
@@ -132,7 +152,7 @@ async def chat(req: Request):
                 log(f"Error Non-Stream: {e}")
                 continue
             finally:
-                ki["in_use"] = False
+                ki["in_use"] = False # SELALU LEPAS KUNCI
                 log(f"RELEASE: key#{ki['index']} (Non-Stream)")
                 
         return JSONResponse({"error": "All keys failed after multiple attempts"}, status_code=500)
@@ -142,13 +162,12 @@ async def chat(req: Request):
     # ==========================================
     async def stream_generator():
         current_body = body.copy()
-        # Deep copy messages agar tidak mengubah aslinya saat di-inject
         current_body["messages"] = [msg.copy() for msg in body.get("messages", [])] 
         
         generated_text_buffer = "" 
         tried_keys = set()
         
-        for attempt in range(len(OLLAMA_KEYS)): # Loop sebanyak jumlah key
+        for attempt in range(len(OLLAMA_KEYS)): 
             if len(tried_keys) >= len(OLLAMA_KEYS):
                 tried_keys.clear()
 
@@ -163,15 +182,17 @@ async def chat(req: Request):
                 if await req.is_disconnected():
                     log("Client membatalkan request saat mengantre stream.")
                     return
-                key = get_available_key(exclude_keys=tried_keys)
-                if key: break
-                await asyncio.sleep(1)
+                
+                # Gunakan fungsi Atomic Lock
+                key = get_and_lock_key(exclude_keys=tried_keys)
+                if key: 
+                    break # Langsung keluar loop, key SUDAH DIKUNCI
+                    
+                await asyncio.sleep(0.5)
 
             ki = key_status[key]
-            ki["in_use"] = True # Kunci key agar tidak dipakai yang lain
             tried_keys.add(key)
-            
-            log(f"STREAM LOCK: key#{ki['index']}")
+            log(f"STREAM LOCK ACQUIRED: key#{ki['index']}")
 
             if generated_text_buffer:
                 log(f"Resuming stream. Injecting {len(generated_text_buffer)} chars.")
@@ -183,7 +204,6 @@ async def chat(req: Request):
                 current_body["messages"] = messages
 
             try:
-                # Timeout read=None memastikan tidak putus otomatis di tengah stream panjang
                 custom_timeout = httpx.Timeout(connect=15.0, read=None, write=15.0, pool=10.0)
                 async with httpx.AsyncClient(timeout=custom_timeout) as client:
                     async with client.stream(
@@ -191,10 +211,16 @@ async def chat(req: Request):
                         json=current_body, headers={"Authorization": f"Bearer {key}"}
                     ) as response:
                         
-                        if response.status_code != 200:
+                        if response.status_code == 429:
                             ki["failures"] += 1
-                            log(f"STREAM ERR {response.status_code}: key#{ki['index']}")
+                            ki["healthy"] = False
+                            log(f"STREAM 429: key#{ki['index']} - Switching key...")
                             continue 
+                            
+                        if response.status_code != 200:
+                            ki["failures"] += 1
+                            log(f"STREAM ERR {response.status_code}: key#{ki['index']} - Switching key...")
+                            continue
 
                         stream_interrupted = False
                         try:
@@ -204,6 +230,7 @@ async def chat(req: Request):
                                         data_str = chunk[6:]
                                         if data_str.strip() == "[DONE]":
                                             ki["success"] += 1
+                                            ki["failures"] = 0
                                             yield chunk + "\n\n"
                                             return 
                                         try:
@@ -231,7 +258,7 @@ async def chat(req: Request):
                 continue
             
             finally:
-                # SANGAT PENTING: Selalu lepas kuncian saat selesai, berhasil, atau gagal
+                # SELALU LEPAS KUNCI
                 ki["in_use"] = False
                 log(f"STREAM RELEASE: key#{ki['index']}")