Spaces:
Running
Running
| import os, httpx | |
| from fastapi import FastAPI, Request | |
| from fastapi.responses import JSONResponse, StreamingResponse, HTMLResponse | |
| import uvicorn | |
| app = FastAPI() | |
| def get_tunnel(): | |
| url = os.environ.get("KAGGLE_TUNNEL_URL", "").rstrip("/") | |
| if not url: | |
| raise ValueError("KAGGLE_TUNNEL_URL not set") | |
| return url | |
| async def root(): | |
| tunnel = os.environ.get("KAGGLE_TUNNEL_URL", "NOT SET") | |
| alive, model = False, "unknown" | |
| if tunnel != "NOT SET": | |
| try: | |
| async with httpx.AsyncClient(timeout=5) as c: | |
| r = await c.get(f"{tunnel}/health") | |
| d = r.json() | |
| alive = d.get("status") == "ok" | |
| model = d.get("model", "unknown") | |
| except: pass | |
| color = "#00ff88" if alive else "#ff4444" | |
| status = "ONLINE" if alive else "OFFLINE" | |
| return HTMLResponse(f"""<html><body style="font-family:monospace;padding:2rem;background:#0d1117;color:#c9d1d9"> | |
| <h2>WirelessAudit Proxy</h2> | |
| <p>Status: <b style="color:{color}">{status}</b></p> | |
| <p>Model: <code>{model}</code></p> | |
| <p>Tunnel: <code>{tunnel}</code></p> | |
| <p>Cline Base URL: <code>/v1</code></p> | |
| </body></html>""") | |
| async def health(): | |
| try: | |
| t = get_tunnel() | |
| async with httpx.AsyncClient(timeout=5) as c: | |
| return (await c.get(f"{t}/health")).json() | |
| except Exception as e: | |
| return JSONResponse({"status": "offline", "error": str(e)}, status_code=503) | |
| async def models(): | |
| try: | |
| t = get_tunnel() | |
| async with httpx.AsyncClient(timeout=10) as c: | |
| return (await c.get(f"{t}/v1/models")).json() | |
| except Exception as e: | |
| return JSONResponse({"error": str(e)}, status_code=503) | |
| async def chat(request: Request): | |
| try: | |
| t = get_tunnel() | |
| except ValueError as e: | |
| return JSONResponse({"error": str(e)}, status_code=503) | |
| body = await request.json() | |
| stream = body.get("stream", False) | |
| try: | |
| if stream: | |
| async def proxy(): | |
| async with httpx.AsyncClient(timeout=120) as c: | |
| async with c.stream("POST", f"{t}/v1/chat/completions", | |
| json=body, headers={"Content-Type": "application/json"}) as r: | |
| async for chunk in r.aiter_bytes(): | |
| yield chunk | |
| return StreamingResponse(proxy(), media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) | |
| async with httpx.AsyncClient(timeout=120) as c: | |
| r = await c.post(f"{t}/v1/chat/completions", json=body, | |
| headers={"Content-Type": "application/json"}) | |
| return JSONResponse(r.json(), status_code=r.status_code) | |
| except httpx.ConnectError: | |
| return JSONResponse({"error": {"message": "Kaggle offline", "code": 503}}, status_code=503) | |
| except Exception as e: | |
| return JSONResponse({"error": {"message": str(e)}}, status_code=500) | |
| async def update_tunnel(request: Request): | |
| body = await request.json() | |
| if body.get("secret") != os.environ.get("HF_TOKEN", ""): | |
| return JSONResponse({"error": "unauthorized"}, status_code=401) | |
| new_url = body.get("tunnel_url", "") | |
| if not new_url: | |
| return JSONResponse({"error": "tunnel_url required"}, status_code=400) | |
| space_id = os.environ.get("MY_SPACE_ID", "") | |
| hf_token = os.environ.get("HF_TOKEN", "") | |
| try: | |
| async with httpx.AsyncClient(timeout=15) as c: | |
| r = await c.post(f"https://huggingface.co/api/spaces/{space_id}/secrets", | |
| headers={"Authorization": f"Bearer {hf_token}", "Content-Type": "application/json"}, | |
| json={"key": "KAGGLE_TUNNEL_URL", "value": new_url}) | |
| if r.status_code == 200: | |
| os.environ["KAGGLE_TUNNEL_URL"] = new_url | |
| return JSONResponse({"status": "updated", "tunnel_url": new_url}) | |
| return JSONResponse({"error": r.text}, status_code=500) | |
| except Exception as e: | |
| return JSONResponse({"error": str(e)}, status_code=500) | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |