Add client auth with ollama-proxy-free

app.py

@@ -2,15 +2,18 @@ import os
 import time
 import random
 import httpx
-from typing import Dict, List
-from fastapi import FastAPI, Request
+from typing import Dict
+from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import JSONResponse, StreamingResponse, Response
 
 app = FastAPI()
 
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
 
-# 3 API KEYS - FAILOVER
+# CLIENT API KEY (for auth)
+MASTER_API_KEY = "ollama-proxy-free"
+
+# BACKEND API KEYS (for Ollama)
 API_KEYS = [
     "8ca25de51e554c099962b78b7ce0c9e9.Mp5dnqctR2zzq3g-NO_M-cjW",
     "dbd1d0c534964684a6d4678348ab8d30.ieDfmSYVnf0MmTjR-AIdNrW9",
@@ -23,8 +26,13 @@ key_status: Dict[str, Dict] = {
     for k in API_KEYS
 }
 
-def pick_key() -> str:
-    """Pick an available key, fallback if all fail"""
+def auth(req: Request):
+    """Validate client API key"""
+    client_key = req.headers.get("Authorization", "").replace("Bearer ", "")
+    if client_key != MASTER_API_KEY:
+        raise HTTPException(401, "Unauthorized")
+
+def pick_key():
     now = time.time()
     valid = [k for k, v in key_status.items() 
              if v["status"] != "dead" and v["cooldown"] < now]
@@ -32,7 +40,7 @@ def pick_key() -> str:
     if valid:
         return random.choice(valid)
     
-    # All failed, reset and try again
+    # Reset if all failed
     for k in key_status:
         key_status[k]["status"] = "active"
         key_status[k]["fail"] = 0
@@ -43,7 +51,6 @@ 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
@@ -54,12 +61,16 @@ def root():
     active = sum(1 for v in key_status.values() if v["status"] == "active")
     return {
         "status": "ok",
+        "mode": "openai-proxy",
         "keys_total": len(API_KEYS),
         "keys_active": active
     }
 
 @app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(req: Request, path: str):
+    # Auth first
+    auth(req)
+    
     target_url = f"{BASE_URL}/v1/{path}"
     body = await req.body()
     
@@ -67,7 +78,7 @@ async def proxy(req: Request, path: str):
     headers.pop("host", None)
     headers.pop("content-length", None)
     
-    # Try each key until success
+    # Try each key
     tried = set()
     for _ in range(len(API_KEYS)):
         key = pick_key()
@@ -89,25 +100,17 @@ async def proxy(req: Request, path: str):
                 
                 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
     
-    # All keys failed
-    return JSONResponse(
-        {"error": "all keys failed", "tried": len(API_KEYS)},
-        status_code=500
-    )
+    return JSONResponse({"error": "all keys failed"}, status_code=500)