Update app.py

app.py

@@ -1,132 +1,241 @@
 import os
 import httpx
 from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse, Response
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.requests import ClientDisconnect
 import time
+import json
+import asyncio
 
 app = FastAPI()
 
+# ==========================================
+# KONFIGURASI
+# ==========================================
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
-MASTER_API_KEY = os.getenv("MASTER_API_KEY", "ollama-proxy-free")
+MASTER_API_KEY = os.getenv("MASTER_API_KEY", "ollama")
 
 OLLAMA_KEYS = []
-for i in range(1, 20):
+for i in range(1, 15):
     key = os.getenv(f"OLLAMA_KEY_{i}")
     if key:
         OLLAMA_KEYS.append(key)
 
+if not OLLAMA_KEYS:
+    OLLAMA_KEYS.append("ollama")
+
+# Status Key
 key_status = {}
 for idx, k in enumerate(OLLAMA_KEYS, 1):
-    key_status[k] = {"index": idx, "prefix": k[:8]+"...", "failures": 0, "success": 0, "last_error": None, "healthy": True}
+    key_status[k] = {
+        "index": idx, 
+        "prefix": k[:8] + "...", 
+        "failures": 0, 
+        "success": 0, 
+        "healthy": True,
+        "in_use": False  # Fitur Lock: 1 Key = 1 Request
+    }
 
 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)"""
+    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
+
+# ==========================================
+# ENDPOINTS
+# ==========================================
 @app.get("/")
 def root():
     return {
         "status": "ok",
-        "base_url": BASE_URL,
-        "master_key": MASTER_API_KEY[:8] + "...",
         "keys_loaded": len(OLLAMA_KEYS),
-        "healthy_keys": sum(1 for v in key_status.values() if v["healthy"]),
-        "keys_status": {v["prefix"]: {"failures": v["failures"], "success": v["success"], "healthy": v["healthy"]} for v in key_status.values()}
+        "keys_status": {
+            v["prefix"]: {
+                "status": "BUSY" if v["in_use"] else "IDLE",
+                "healthy": v["healthy"],
+                "success": v["success"]
+            } for v in key_status.values()
+        }
     }
 
-def get_healthy_keys(max_failures=2):
-    """Get keys that are healthy (low failures)"""
-    healthy = [k for k, v in key_status.items() if v["failures"] < max_failures and v["healthy"]]
-    if not healthy:
-        for v in key_status.values():
-            v["failures"] = 0
-            v["healthy"] = True
-        return OLLAMA_KEYS[:2]
-    return healthy[:2]
-
 @app.post("/v1/chat/completions")
 async def chat(req: Request):
     auth_key = req.headers.get("Authorization", "").replace("Bearer ", "")
     if auth_key != MASTER_API_KEY:
-        log(f"AUTH FAIL: received '{auth_key[:8]}...' expected '{MASTER_API_KEY[:8]}...'")
         return JSONResponse({"error": "Unauthorized"}, status_code=401)
     
-    body = await req.json()
+    # Tangkap error jika client kabur sebelum request terkirim penuh
+    try:
+        body = await req.json()
+    except ClientDisconnect:
+        log("Client kabur sebelum proxy selesai membaca request body.")
+        return Response(status_code=499)
+    except json.JSONDecodeError:
+        return JSONResponse({"error": "Invalid JSON body"}, status_code=400)
+
     is_stream = body.get("stream", False)
-    model = body.get("model", "?")
-    log(f"REQUEST: model='{model}'")
-    
-    candidate_keys = get_healthy_keys(max_failures=2)
-    log(f"Using top {len(candidate_keys)} keys")
-    
-    for attempt, key in enumerate(candidate_keys):
-        ki = key_status[key]
-        log(f"TRY #{attempt+1}: using key#{ki['index']} ({ki['prefix']})")
-        
-        try:
-            start_time = time.time()
-            async with httpx.AsyncClient(timeout=30.0) as client:
-                resp = await client.post(f"{BASE_URL}/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {key}"})
-            elapsed = time.time() - start_time
+
+    # ==========================================
+    # LOGIKA NON-STREAM 
+    # ==========================================
+    if not is_stream:
+        tried_keys = set()
+        for attempt in range(len(OLLAMA_KEYS)): # Loop sebanyak jumlah key
+            if len(tried_keys) >= len(OLLAMA_KEYS):
+                tried_keys.clear()
+
+            key = None
+            log("Menunggu API Key idle (Antrean Non-Stream)...")
             
-            if resp.status_code == 200:
-                ki["success"] += 1
-                ki["failures"] = 0
-                ki["healthy"] = True
-                log(f"SUCCESS: key#{ki['index']} responded in {elapsed:.2f}s")
-                return Response(resp.content, status_code=200)
+            # Antrean Tanpa Batas Waktu
+            while True:
+                if await req.is_disconnected():
+                    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)
+
+            ki = key_status[key]
+            ki["in_use"] = True
+            tried_keys.add(key)
+            log(f"LOCK: key#{ki['index']} (Non-Stream)")
             
-            elif resp.status_code == 429:
+            try:
+                async with httpx.AsyncClient(timeout=120.0) as client:
+                    resp = await client.post(
+                        f"{BASE_URL}/v1/chat/completions",
+                        json=body,
+                        headers={"Authorization": f"Bearer {key}"}
+                    )
+                if resp.status_code == 200:
+                    ki["success"] += 1
+                    return Response(content=resp.content, media_type=resp.headers.get("content-type"))
+                else:
+                    ki["failures"] += 1
+                    continue
+            except Exception as e:
                 ki["failures"] += 1
-                ki["healthy"] = False
-                log(f"RATE LIMIT: key#{ki['index']} - skip to next")
+                log(f"Error Non-Stream: {e}")
                 continue
+            finally:
+                ki["in_use"] = False
+                log(f"RELEASE: key#{ki['index']} (Non-Stream)")
+                
+        return JSONResponse({"error": "All keys failed after multiple attempts"}, status_code=500)
+
+    # ==========================================
+    # LOGIKA STREAMING (Seamless Fallback + Queue)
+    # ==========================================
+    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
+            if len(tried_keys) >= len(OLLAMA_KEYS):
+                tried_keys.clear()
+
+            key = None
+            if attempt == 0:
+                log("Menunggu API Key idle (Antrean Stream Baru)...")
+            else:
+                log(f"Menunggu API Key idle (Antrean Fallback ke-{attempt})...")
+                
+            # Antrean Tanpa Batas Waktu
+            while True:
+                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)
+
+            ki = key_status[key]
+            ki["in_use"] = True # Kunci key agar tidak dipakai yang lain
+            tried_keys.add(key)
             
-            elif resp.status_code >= 500:
+            log(f"STREAM LOCK: key#{ki['index']}")
+
+            if generated_text_buffer:
+                log(f"Resuming stream. Injecting {len(generated_text_buffer)} chars.")
+                messages = current_body.get("messages", [])
+                if messages and messages[-1].get("role") == "assistant":
+                    messages[-1]["content"] = generated_text_buffer
+                else:
+                    messages.append({"role": "assistant", "content": generated_text_buffer})
+                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(
+                        "POST", f"{BASE_URL}/v1/chat/completions",
+                        json=current_body, headers={"Authorization": f"Bearer {key}"}
+                    ) as response:
+                        
+                        if response.status_code != 200:
+                            ki["failures"] += 1
+                            log(f"STREAM ERR {response.status_code}: key#{ki['index']}")
+                            continue 
+
+                        stream_interrupted = False
+                        try:
+                            async for chunk in response.aiter_lines():
+                                if chunk:
+                                    if chunk.startswith("data: "):
+                                        data_str = chunk[6:]
+                                        if data_str.strip() == "[DONE]":
+                                            ki["success"] += 1
+                                            yield chunk + "\n\n"
+                                            return 
+                                        try:
+                                            data_json = json.loads(data_str)
+                                            if "choices" in data_json and len(data_json["choices"]) > 0:
+                                                delta = data_json["choices"][0].get("delta", {})
+                                                content = delta.get("content", "")
+                                                if content:
+                                                    generated_text_buffer += content
+                                        except json.JSONDecodeError:
+                                            pass
+                                    yield chunk + "\n\n"
+                                    
+                        except (httpx.ReadTimeout, httpx.ReadError, httpx.RemoteProtocolError) as e:
+                            log(f"STREAM PUTUS: key#{ki['index']}. Buffering...")
+                            ki["failures"] += 1
+                            stream_interrupted = True
+                        
+                        if not stream_interrupted:
+                            return 
+
+            except Exception as e:
                 ki["failures"] += 1
-                ki["last_error"] = f"http {resp.status_code}"
-                log(f"SERVER ERROR: key#{ki['index']} got {resp.status_code}, trying next")
+                log(f"STREAM EXCEPTION: key#{ki['index']} - {e}")
                 continue
             
-            else:
-                ki["last_error"] = f"http {resp.status_code}"
-                log(f"ERROR: key#{ki['index']} got {resp.status_code}")
-                return Response(resp.content, status_code=resp.status_code)
+            finally:
+                # SANGAT PENTING: Selalu lepas kuncian saat selesai, berhasil, atau gagal
+                ki["in_use"] = False
+                log(f"STREAM RELEASE: key#{ki['index']}")
                 
-        except httpx.TimeoutException:
-            ki["failures"] += 1
-            ki["healthy"] = False
-            ki["last_error"] = "timeout after 30s"
-            log(f"TIMEOUT: key#{ki['index']} - already healthy=False, try next")
-            continue
-            
-        except Exception as e:
-            ki["failures"] += 1
-            ki["last_error"] = str(e)[:50]
-            log(f"EXCEPTION: key#{ki['index']} error: {e}")
-            continue
-    
-    log(f"ALL KEYS FAILED for model='{model}'")
-    return JSONResponse({"error": "all keys failed after 2 attempts", "model": model,
-        "keys_status": {v["prefix"]: {"failures": v["failures"], "last_error": v["last_error"]} for v in key_status.values()}}, status_code=500)
+        yield f"data: {json.dumps({'error': 'Stream failed completely'})}\n\ndata: [DONE]\n\n"
 
-@app.get("/v1/models")
-def models(req: Request):
-    auth_key = req.headers.get("Authorization", "").replace("Bearer ", "")
-    if auth_key != MASTER_API_KEY:
-        return JSONResponse({"error": "Unauthorized"}, status_code=401)
-    
-    log("REQUEST: GET /v1/models")
-    
-    for key in get_healthy_keys(max_failures=2):
-        ki = key_status[key]
-        try:
-            resp = httpx.get(f"{BASE_URL}/v1/models", headers={"Authorization": f"Bearer {key}"}, timeout=10)
-            if resp.status_code == 200:
-                ki["success"] += 1
-                log(f"MODELS OK: key#{ki['index']}")
-                return Response(resp.content, status_code=200)
-        except Exception as e:
-            ki["last_error"] = str(e)[:50]
-            log(f"MODELS FAIL: key#{ki['index']} - {e}")
-    
-    return JSONResponse({"error": "no keys available"}, status_code=500)
+    return StreamingResponse(stream_generator(), media_type="text/event-stream")
+