import os import secrets from datetime import datetime, timezone from fastapi import FastAPI, HTTPException, Header, status, APIRouter from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel import redis.asyncio as aioredis import httpx app = FastAPI(title="Zero-Cost API Key Orchestrator") # Enable CORS for cross-origin frontend dashboards app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"] ) redis_conn = aioredis.from_url("redis://localhost:6379", decode_responses=True) class RouterPayload(BaseModel): model: str fallback_index: int = 0 class SSOLoginPayload(BaseModel): provider: str access_token: str class ReportFailurePayload(BaseModel): model: str provider: str endpoint: str DAILY_TEXT_LIMIT = 100 DAILY_IMAGE_LIMIT = 10 PROVIDER_ROUTING_MATRIX = { "meta-llama/llama-3.3-70b-instruct": [{"endpoint": "https://api-inference.huggingface.co/models/meta-llama/Llama-3.3-70B-Instruct", "provider": "huggingface"}], "black-forest-labs/flux.1-schnell": [{"endpoint": "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell", "provider": "huggingface"}] } MODEL_LIMITS = {"text": 120, "image": 15} # FIXED: Added root endpoint to eliminate the 404 "Not Found" error on page load @app.get("/") async def root_health_check(): return { "status": "online", "message": "Zero-Cost API Key Orchestrator is running perfectly.", "endpoints": { "orchestrate": "/v1/orchestrate [POST]", "sso_callback": "/v1/auth/sso-callback [POST]", "report_failure": "/v1/report-failure [POST]" } } @app.post("/v1/auth/sso-callback") async def sso_callback(payload: SSOLoginPayload): async with httpx.AsyncClient() as client: try: if payload.provider == "google": res = await client.get( "https://www.googleapis.com/oauth2/v3/userinfo", headers={"Authorization": f"Bearer {payload.access_token}"} ) if res.status_code != 200: raise HTTPException(status_code=401, detail="Google token validation failed.") user_info = res.json() unique_id = f"google_{user_info['sub']}" email = user_info["email"] elif payload.provider == "github": res = await client.get( "https://api.github.com/user", headers={"Authorization": f"Bearer {payload.access_token}"} ) if res.status_code != 200: raise HTTPException(status_code=401, detail="GitHub token validation failed.") user_info = res.json() unique_id = f"github_{user_info['id']}" email = user_info.get("email") or f"{unique_id}@platform.internal" else: raise HTTPException(status_code=400, detail="Invalid provider entry.") except Exception as e: raise HTTPException(status_code=401, detail=f"OAuth verification failure: {str(e)}") user_token_lookup_key = f"user:token:{unique_id}" existing_token = await redis_conn.get(user_token_lookup_key) if existing_token: active_platform_key = existing_token else: active_platform_key = f"sk_platform_{secrets.token_hex(24)}" await redis_conn.set(f"user:profile:{unique_id}", email) await redis_conn.set(f"token_map:{active_platform_key}", unique_id) await redis_conn.set(user_token_lookup_key, active_platform_key) await redis_conn.set(f"quota:{unique_id}:text", DAILY_TEXT_LIMIT) await redis_conn.set(f"quota:{unique_id}:image", DAILY_IMAGE_LIMIT) return {"status": "authenticated", "platform_bearer_token": active_platform_key} @app.post("/v1/orchestrate") async def orchestrate_routing(payload: RouterPayload, authorization: str = Header(None)): if not authorization or not authorization.startswith("Bearer "): raise HTTPException(status_code=401, detail="Missing or invalid token layout.") platform_key = authorization.split(" ")[1] user_id = await redis_conn.get(f"token_map:{platform_key}") if not user_id: raise HTTPException(status_code=401, detail="Token unregistered or expired.") category = "image" if "flux" in payload.model.lower() else "text" rpm_limit = MODEL_LIMITS.get(category, 60) global_breaker_key = "platform:global:concurrency" async with redis_conn.pipeline(transaction=True) as pipe: pipe.incr(global_breaker_key) pipe.expire(global_breaker_key, 1) global_res = await pipe.execute() if global_res[0] > 50000: raise HTTPException(status_code=503, detail="Platform Overload.") user_model_rate_key = f"key:{platform_key}:model:{payload.model}:rpm" async with redis_conn.pipeline(transaction=True) as pipe: pipe.incr(user_model_rate_key) pipe.ttl(user_model_rate_key) current_count, current_ttl = await pipe.execute() if current_ttl == -1: await redis_conn.expire(user_model_rate_key, 60) if current_count > rpm_limit: raise HTTPException(status_code=429, detail="Per-model request limit reached.") today_str = datetime.now(timezone.utc).strftime("%Y-%m-%d") daily_reset_key = f"user:{user_id}:reset:{today_str}" if await redis_conn.setnx(daily_reset_key, "checked"): await redis_conn.set(f"quota:{user_id}:text", DAILY_TEXT_LIMIT) await redis_conn.set(f"quota:{user_id}:image", DAILY_IMAGE_LIMIT) await redis_conn.expire(daily_reset_key, 86400) quota_key = f"quota:{user_id}:{category}" balance = await redis_conn.get(quota_key) if not balance or int(balance) <= 0: raise HTTPException(status_code=429, detail=f"Daily credit allocation for {category} models exhausted.") await redis_conn.decrby(quota_key, 1) endpoints = PROVIDER_ROUTING_MATRIX.get( payload.model.lower(), [{"endpoint": f"https://api-inference.huggingface.co/models/{payload.model}", "provider": "huggingface"}] ) for ep in endpoints: bl_key = f"blacklist:model:{payload.model.lower()}:provider:{ep['provider']}:endpoint:{ep['endpoint']}" if not await redis_conn.get(bl_key): return { "status": "approved", "target_endpoint": ep["endpoint"], "provider_prefix": ep["provider"], "data_category": category } raise HTTPException(status_code=502, detail="All endpoints temporarily down.") @app.post("/v1/report-failure") async def report_failure(payload: ReportFailurePayload): bl_key = f"blacklist:model:{payload.model.lower()}:provider:{payload.provider}:endpoint:{payload.endpoint}" await redis_conn.set(bl_key, "dead", ex=300) return {"status": "reported"}