Spaces:
Sleeping
Sleeping
| import itertools | |
| from fastapi import FastAPI, Request, Response, UploadFile, File, Form | |
| from fastapi.responses import StreamingResponse, JSONResponse | |
| from starlette.background import BackgroundTask | |
| import httpx | |
| import uvicorn | |
| import os | |
| import io | |
| import uuid | |
| import asyncio | |
| app = FastAPI(title="Featherless Proxy") | |
| # --- Featherless API Keys --- | |
| api_keys_env = os.environ.get("FEATHERLESS_API_KEYS", "") | |
| if api_keys_env: | |
| # Support comma-separated or newline-separated keys in env var | |
| api_keys = [k.strip() for k in api_keys_env.replace('\\n', ',').replace('\n', ',').split(",") if k.strip()] | |
| else: | |
| API_KEYS_FILE = "api_keys.txt" | |
| if not os.path.exists(API_KEYS_FILE): | |
| API_KEYS_FILE = os.path.join(os.path.dirname(__file__), "api_keys.txt") | |
| if os.path.exists(API_KEYS_FILE): | |
| with open(API_KEYS_FILE, "r") as f: | |
| api_keys = [line.strip() for line in f if line.strip()] | |
| else: | |
| api_keys = [] | |
| if not api_keys: | |
| raise ValueError("No API keys found. Please set the FEATHERLESS_API_KEYS environment variable.") | |
| key_pool = itertools.cycle(api_keys) | |
| TARGET_BASE_URL = "https://api.featherless.ai" | |
| client = httpx.AsyncClient(base_url=TARGET_BASE_URL, timeout=None) | |
| # --- NIM API Keys --- | |
| nim_keys_env = os.environ.get("NIM_API_KEYS", "") | |
| if nim_keys_env: | |
| nim_keys = [k.strip() for k in nim_keys_env.replace('\\n', ',').replace('\n', ',').split(",") if k.strip()] | |
| else: | |
| NIM_KEYS_FILE = os.path.join(os.path.dirname(__file__), "nim_keys.txt") | |
| nim_keys = [] | |
| if os.path.exists(NIM_KEYS_FILE): | |
| with open(NIM_KEYS_FILE, "r") as f: | |
| nim_keys = [line.strip() for line in f if line.strip()] | |
| nim_key_pool = itertools.cycle(nim_keys) if nim_keys else None | |
| NIM_BASE_URL = "https://integrate.api.nvidia.com" | |
| nim_client = httpx.AsyncClient(base_url=NIM_BASE_URL, timeout=120.0) | |
| async def shutdown_event(): | |
| await client.aclose() | |
| if nim_client: | |
| await nim_client.aclose() | |
| async def health_check(): | |
| return {"status": "ok"} | |
| async def analyze_content( | |
| text: str = Form(...) | |
| ): | |
| if not nim_key_pool: | |
| return JSONResponse(status_code=500, content={"error": "No NIM keys available"}) | |
| next_key = next(nim_key_pool) | |
| headers = { | |
| "Authorization": f"Bearer {next_key}", | |
| "Content-Type": "application/json" | |
| } | |
| messages = [] | |
| try: | |
| model = "deepseek-ai/deepseek-v4-flash-0731" | |
| if not text: | |
| return JSONResponse(status_code=400, content={"error": "Must provide text."}) | |
| messages.append({"role": "user", "content": text}) | |
| payload = { | |
| "model": model, | |
| "messages": messages, | |
| "max_tokens": 4096 | |
| } | |
| try: | |
| resp = await nim_client.post("/v1/chat/completions", json=payload, headers=headers, timeout=60.0) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| except Exception as primary_e: | |
| fallback_model = "meta/llama-3.1-8b-instruct" | |
| print(f"Primary model {model} failed/timed out ({type(primary_e).__name__}). Falling back to {fallback_model}...") | |
| payload["model"] = fallback_model | |
| resp = await nim_client.post("/v1/chat/completions", json=payload, headers=headers, timeout=30.0) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| return JSONResponse(content={"result": data["choices"][0]["message"]["content"]}) | |
| except Exception as e: | |
| err_msg = str(e) or repr(e) | |
| print(f"Analyze error: {err_msg}") | |
| return JSONResponse(status_code=500, content={"error": err_msg}) | |
| async def proxy(path: str, request: Request): | |
| next_key = next(key_pool) | |
| print(f"Using API Key: {next_key[:8]}... for path: {path}") | |
| # Forward headers, but replace Authorization with the selected key | |
| headers = dict(request.headers) | |
| headers.pop("host", None) # Let httpx handle host | |
| headers["authorization"] = f"Bearer {next_key}" | |
| url = f"/{path}" | |
| # Forward query params | |
| params = dict(request.query_params) | |
| # Read the body if any | |
| body = await request.body() | |
| req = client.build_request( | |
| method=request.method, | |
| url=url, | |
| content=body, | |
| headers=headers, | |
| params=params, | |
| ) | |
| try: | |
| response = await client.send(req, stream=True) | |
| response_headers = dict(response.headers) | |
| # Remove encodings like gzip/deflate if we are proxying, | |
| # or just pass through. It's safer to delete 'content-encoding' and 'content-length' | |
| # when streaming because Starlette will set them. | |
| response_headers.pop("content-encoding", None) | |
| response_headers.pop("content-length", None) | |
| return StreamingResponse( | |
| response.aiter_bytes(), | |
| status_code=response.status_code, | |
| headers=response_headers, | |
| background=BackgroundTask(response.aclose) | |
| ) | |
| except Exception as e: | |
| print(f"Proxy error: {e}") | |
| return Response(content=f"Proxy error: {str(e)}", status_code=500) | |
| if __name__ == "__main__": | |
| uvicorn.run("server:app", host="0.0.0.0", port=8000) | |