Spaces:
Running
Running
File size: 4,362 Bytes
109c2ac 8464010 109c2ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | 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
@app.get("/")
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>""")
@app.get("/health")
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)
@app.get("/v1/models")
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)
@app.post("/v1/chat/completions")
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)
@app.post("/update-tunnel")
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)
|