Update app.py

app.py

@@ -5,7 +5,7 @@ import httpx
 from typing import Dict
 
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import Response, StreamingResponse, JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 
 app = FastAPI()
 
@@ -26,25 +26,23 @@ def load_keys(prefix="OLLAMA_KEY_"):
     return keys
 
 API_KEYS = load_keys()
-print(f"[INIT] Loaded {len(API_KEYS)} keys")
+print(f"[INIT] Loaded {len(API_KEYS)} keys: {[k[:10] for k in API_KEYS}")
 
 # ─────────────────────────────
 # 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
+# AUTH - Accept valid backend key OR master key
 # ─────────────────────────────
 def auth(req: Request):
-    if req.headers.get("Authorization") != f"Bearer {MASTER_API_KEY}":
+    auth_header = req.headers.get("Authorization", "").replace("Bearer ", "")
+    # Accept if: (1) it's a valid backend key OR (2) it's master key
+    if auth_header not in API_KEYS and auth_header != MASTER_API_KEY:
         raise HTTPException(401, "Unauthorized")
 
 # ─────────────────────────────
@@ -89,16 +87,13 @@ async def proxy(req: Request, path: str):
     for _ in range(len(API_KEYS)):
         key = pick_key()
         if not key:
-            return JSONResponse({"error": "all keys exhausted"}, 503)
+            return JSONResponse({"error": "all keys failed"}, status_code=503)
 
         headers["Authorization"] = f"Bearer {key}"
 
         try:
             async with httpx.AsyncClient(timeout=60) as client:
-
-                # ─────────────────────────────
-                # STEP 1: VALIDATE REQUEST (NO STREAM)
-                # ─────────────────────────────
+                # Validate request first
                 test = await client.request(
                     method=req.method,
                     url=target_url,
@@ -114,13 +109,18 @@ async def proxy(req: Request, path: str):
                         key_status[key]["fail"] += 1
                         if key_status[key]["fail"] >= 3:
                             mark_dead(key)
-                    continue
+                        continue
+                    
+                    mark_ok(key)
+                    return Response(
+                        content=test.content,
+                        status_code=test.status_code,
+                        headers=dict(test.headers)
+                    )
 
                 mark_ok(key)
 
-                # ─────────────────────────────
-                # STEP 2: STREAM RESPONSE (SAFE)
-                # ─────────────────────────────
+                # Stream response
                 stream = await client.stream(
                     method=req.method,
                     url=target_url,
@@ -131,26 +131,19 @@ async def proxy(req: Request, path: str):
 
                 return StreamingResponse(
                     stream.aiter_bytes(),
-                    media_type=test.headers.get(
-                        "content-type",
-                        "application/json"
-                    )
+                    media_type=test.headers.get("content-type", "application/json")
                 )
 
         except Exception as e:
             key_status[key]["fail"] += 1
-
             if key_status[key]["fail"] >= 3:
                 mark_dead(key)
-
-            print(f"[ERROR] key={key[:6]} -> {e}")
+            print(f"[ERROR] key={key[:10]} -> {e}")
+            continue
 
     return JSONResponse(
-        {
-            "error": "all keys failed",
-            "hint": "check keys / model access / BASE_URL"
-        },
-        500
+        {"error": "all keys failed", "hint": "check keys / model access / BASE_URL"},
+        status_code=500
     )
 
 # ─────────────────────────────
@@ -158,8 +151,10 @@ async def proxy(req: Request, path: str):
 # ─────────────────────────────
 @app.get("/")
 def root():
+    active = sum(1 for v in key_status.values() if v["status"] == "active")
     return {
         "status": "ok",
         "mode": "openai-proxy-fixed",
-        "keys": len(API_KEYS)
+        "keys": len(API_KEYS),
+        "active": active
     }