""" NVIDIA NIM Proxy — Gradio/FastAPI Edition Runs FOR FREE on Hugging Face Spaces (CPU Basic, no credit card needed). """ import asyncio import json import os import re import time from typing import Optional, Union import httpx import gradio as gr from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse # ── Config ────────────────────────────────────────────────────────────────── NVIDIA_URL = os.environ.get("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1") COOLDOWN = int(os.environ.get("KEY_COOLDOWN_SECONDS", str(15 * 60))) KEEPALIVE = float(os.environ.get("KEEPALIVE_INTERVAL", "5")) RETRIES = int(os.environ.get("TRANSIENT_RETRIES", "3")) BACKOFF = float(os.environ.get("BACKOFF_BASE", "1.0")) PROXY_AUTH_TOKEN = os.environ.get("PROXY_AUTH_TOKEN", "").strip() THINK_MIN_TOKENS = int(os.environ.get("THINK_MIN_TOKENS", "202000")) ALLOWED = { "model", "messages", "tools", "tool_choice", "temperature", "top_p", "top_k", "max_tokens", "stream", "seed", "stop", "response_format", "frequency_penalty", "presence_penalty", "logprobs", "top_logprobs", "n", "chat_template_kwargs", } KEYS: list[dict] = [] for i in range(1, 9): k = os.environ.get(f"NVIDIA_KEY_{i}") if k and k.strip(): KEYS.append({"i": i, "key": k.strip()}) cooldown: dict[int, float] = {} _rr = 0 # ── FastAPI App ──────────────────────────────────────────────────────────── app = FastAPI() def _err(code: int, msg: str) -> JSONResponse: return JSONResponse(status_code=code, content={"error": {"message": msg, "type": "error"}}) def _check_auth(auth: Optional[str]) -> Optional[JSONResponse]: if not PROXY_AUTH_TOKEN: return _err(401, "PROXY_AUTH_TOKEN not configured.") if not auth: return _err(401, "Missing Authorization header.") parts = auth.split(" ", 1) if len(parts) != 2 or parts[0].lower() != "bearer": return _err(401, "Malformed Authorization header.") if parts[1].strip() != PROXY_AUTH_TOKEN: return _err(401, "Invalid token.") return None def _build_payload(body: dict) -> dict: extra = body.pop("extra_body", None) if isinstance(extra, dict): body.update(extra) if "max_completion_tokens" in body and "max_tokens" not in body: body["max_tokens"] = body["max_completion_tokens"] body = {k: v for k, v in body.items() if k in ALLOWED} body.pop("reasoning_effort", None) model = str(body.get("model", "")).lower() if "glm" in model or "gemma" in model: ctk = body.get("chat_template_kwargs") if isinstance(ctk, dict): ctk.setdefault("enable_thinking", True) ctk.setdefault("reasoning_effort", "max") ctk.pop("max_thinking_tokens", None) body["chat_template_kwargs"] = ctk else: body["chat_template_kwargs"] = {"enable_thinking": True, "reasoning_effort": "max"} mt = body.get("max_tokens") try: mt = int(mt) except: mt = 0 if mt < THINK_MIN_TOKENS: body["max_tokens"] = THINK_MIN_TOKENS return body async def _post(url, payload, accept="text/event-stream"): if not KEYS: return _err(503, "No NVIDIA keys configured.") global _rr last_status, last_detail = None, "No response." for attempt in range(RETRIES + 1): order = [KEYS[(_rr + i) % len(KEYS)] for i in range(len(KEYS))] _rr = (_rr + 1) % len(KEYS) now = time.time() avail = [e for e in order if now >= cooldown.get(e["i"], 0)] if not avail: return _err(429, f"All keys cooling down ({COOLDOWN // 60} min). Retry later.") transient = False for e in avail: headers = {"Authorization": f"Bearer {e['key']}", "Content-Type": "application/json", "Accept": accept, "Accept-Encoding": "identity"} try: async with httpx.AsyncClient() as client: resp = await client.post(url, json=payload, headers=headers) except Exception as ex: last_status, last_detail, transient = 502, str(ex), True continue status = resp.status_code if status == 429 or status >= 500: if status == 429: cooldown[e["i"]] = time.time() + COOLDOWN last_detail, last_status, transient = resp.text[:2000], status, True continue elif status >= 400: detail = resp.text[:2000] if "DEGRADED" in detail.upper(): last_status, last_detail, transient = status, detail, True continue return _err(status, detail) else: return resp if transient and attempt < RETRIES: await asyncio.sleep(BACKOFF * (attempt + 1)) return _err(last_status or 502, f"All keys failed after {RETRIES + 1} tries. Last: HTTP {last_status} - {last_detail}") async def _forward(resp: httpx.Response): async for chunk in resp.aiter_raw(): yield chunk @app.post("/v1/chat/completions") async def chat(request: Request): auth_err = _check_auth(request.headers.get("authorization")) if auth_err: return auth_err try: body = await request.json() except: return _err(400, "Invalid JSON.") if not isinstance(body, dict): return _err(400, "Body must be a JSON object.") payload = _build_payload(body) is_stream = bool(payload.get("stream", False)) if is_stream: async def stream_with_keepalive(): yield b": proxy-connected\n\n" result = await _post(f"{NVIDIA_URL}/chat/completions", payload) if isinstance(result, JSONResponse): yield b"data: " + result.body + b"\n\n" return async for chunk in _forward(result): yield chunk return StreamingResponse(stream_with_keepalive(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) result = await _post(f"{NVIDIA_URL}/chat/completions", payload) if isinstance(result, JSONResponse): return result return Response(result.content, media_type=result.headers.get("content-type", "application/json")) @app.get("/health") async def health(): return {"status": "ok", "keys": len(KEYS)} @app.get("/models") async def models(): return {"data": [{"id": "z-ai/glm-5.2", "object": "model"}, {"id": "google/diffusiongemma-26b-a4b-it", "object": "model"}]} # ── Gradio UI Wrapper ────────────────────────────────────────────────────── with gr.Blocks(title="NIM Proxy") as demo: gr.Markdown("# ✅ NVIDIA NIM Proxy is Running") gr.Markdown("API is accessible at `/v1/chat/completions`") app = gr.mount_gradio_app(app, demo, path="/") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)