Add 3-key failover

app.py

@@ -1,22 +1,65 @@
 import os
+import time
+import random
 import httpx
-import traceback
+from typing import Dict, List
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, StreamingResponse, Response
 
 app = FastAPI()
 
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
-TEST_KEY = "37e81a6be4104fbfbfbe2ecf557a2c10.GoIbzpHebdM9ZcHarQ9A12Cp"
+
+# 3 API KEYS - FAILOVER
+API_KEYS = [
+    "8ca25de51e554c099962b78b7ce0c9e9.Mp5dnqctR2zzq3g-NO_M-cjW",
+    "dbd1d0c534964684a6d4678348ab8d30.ieDfmSYVnf0MmTjR-AIdNrW9",
+    "37e81a6be4104fbfbfbe2ecf557a2c10.GoIbzpHebdM9ZcHarQ9A12Cp"
+]
+
+# KEY STATUS
+key_status: Dict[str, Dict] = {
+    k: {"fail": 0, "cooldown": 0, "status": "active"}
+    for k in API_KEYS
+}
+
+def pick_key() -> str:
+    """Pick an available key, fallback if all fail"""
+    now = time.time()
+    valid = [k for k, v in key_status.items() 
+             if v["status"] != "dead" and v["cooldown"] < now]
+    
+    if valid:
+        return random.choice(valid)
+    
+    # All failed, reset and try again
+    for k in key_status:
+        key_status[k]["status"] = "active"
+        key_status[k]["fail"] = 0
+    
+    return random.choice(API_KEYS)
+
+def mark_fail(key: str):
+    key_status[key]["fail"] += 1
+    if key_status[key]["fail"] >= 3:
+        key_status[key]["cooldown"] = time.time() + 60
+        print(f"[KEY] {key[:8]}... -> cooldown 60s")
+
+def mark_ok(key: str):
+    key_status[key]["fail"] = 0
+    key_status[key]["status"] = "active"
 
 @app.get("/")
 def root():
-    return {"status": "ok", "base_url": BASE_URL}
+    active = sum(1 for v in key_status.values() if v["status"] == "active")
+    return {
+        "status": "ok",
+        "keys_total": len(API_KEYS),
+        "keys_active": active
+    }
 
-@app.api_route("/v1/{path:path}", methods=["GET", "POST"])
+@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(req: Request, path: str):
-    print(f"[PROXY] path={path}")
-    
     target_url = f"{BASE_URL}/v1/{path}"
     body = await req.body()
     
@@ -24,29 +67,47 @@ async def proxy(req: Request, path: str):
     headers.pop("host", None)
     headers.pop("content-length", None)
     
-    # Hardcoded key
-    headers["Authorization"] = f"Bearer {TEST_KEY}"
-    
-    print(f"[PROXY] target={target_url}")
-    print(f"[PROXY] key={TEST_KEY[:10]}...")
+    # Try each key until success
+    tried = set()
+    for _ in range(len(API_KEYS)):
+        key = pick_key()
+        if key in tried:
+            continue
+        tried.add(key)
+        
+        headers["Authorization"] = f"Bearer {key}"
+        
+        try:
+            async with httpx.AsyncClient(timeout=60) as client:
+                resp = await client.request(
+                    method=req.method,
+                    url=target_url,
+                    headers=headers,
+                    content=body,
+                    params=req.query_params
+                )
+                
+                if resp.status_code < 400:
+                    mark_ok(key)
+                    print(f"[OK] key={key[:6]}...")
+                    return Response(
+                        content=resp.content,
+                        status_code=resp.status_code,
+                        headers=dict(resp.headers)
+                    )
+                
+                # Error with this key
+                mark_fail(key)
+                print(f"[FAIL] key={key[:6]}... status={resp.status_code}")
+                continue
+                
+        except Exception as e:
+            mark_fail(key)
+            print(f"[ERROR] key={key[:6]}... {e}")
+            continue
     
-    try:
-        async with httpx.AsyncClient(timeout=60) as client:
-            resp = await client.request(
-                method=req.method,
-                url=target_url,
-                headers=headers,
-                content=body,
-                params=req.query_params
-            )
-            print(f"[PROXY] status={resp.status_code}")
-            
-            return Response(
-                content=resp.content,
-                status_code=resp.status_code,
-                headers=dict(resp.headers)
-            )
-    except Exception as e:
-        print(f"[ERROR] {e}")
-        traceback.print_exc()
-        return JSONResponse({"error": str(e)}, status_code=500)
+    # All keys failed
+    return JSONResponse(
+        {"error": "all keys failed", "tried": len(API_KEYS)},
+        status_code=500
+    )