Spaces:
Sleeping
Sleeping
| import os | |
| import time | |
| import random | |
| import httpx | |
| 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") | |
| # 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", | |
| "37e81a6be4104fbfbfbe2ecf557a2c10.GoIbzpHebdM9ZcHarQ9A12Cp" | |
| ] | |
| # KEY STATUS | |
| key_status: Dict[str, Dict] = { | |
| k: {"fail": 0, "cooldown": 0, "status": "active"} | |
| for k in API_KEYS | |
| } | |
| 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] | |
| if valid: | |
| return random.choice(valid) | |
| # Reset if all failed | |
| 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 | |
| def mark_ok(key: str): | |
| key_status[key]["fail"] = 0 | |
| key_status[key]["status"] = "active" | |
| 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 | |
| } | |
| async def proxy(req: Request, path: str): | |
| # Auth first | |
| auth(req) | |
| target_url = f"{BASE_URL}/v1/{path}" | |
| body = await req.body() | |
| headers = dict(req.headers) | |
| headers.pop("host", None) | |
| headers.pop("content-length", None) | |
| # Try each key | |
| 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) | |
| return Response( | |
| content=resp.content, | |
| status_code=resp.status_code, | |
| headers=dict(resp.headers) | |
| ) | |
| mark_fail(key) | |
| continue | |
| except Exception as e: | |
| mark_fail(key) | |
| continue | |
| return JSONResponse({"error": "all keys failed"}, status_code=500) | |
| # rebuild | |
| # Testing build | |