Spaces:
Running
Running
| import os, json, time, tempfile, httpx | |
| from fastapi import FastAPI, Request | |
| from fastapi.responses import JSONResponse, StreamingResponse, HTMLResponse | |
| import uvicorn | |
| app = FastAPI() | |
| KIMI_BASE = os.environ.get("KIMI_BASE_URL", "https://api.kimi.com/coding").rstrip("/") | |
| DEFAULT_MODEL = os.environ.get("KIMI_DEFAULT_MODEL", "k3") | |
| DAILY_CAP = int(os.environ.get("DAILY_TOKEN_CAP", "2500000")) | |
| USAGE_FILE = os.path.join(tempfile.gettempdir(), "wirelessaudit_token_usage.json") | |
| _usage = {"date": "", "used": 0} | |
| def _today(): | |
| return time.strftime("%Y-%m-%d", time.gmtime()) | |
| def _load_usage(): | |
| try: | |
| with open(USAGE_FILE) as f: | |
| d = json.load(f) | |
| if d.get("date") == _today() and isinstance(d.get("used"), int): | |
| _usage.update(d) | |
| except Exception: | |
| pass | |
| def _save_usage(): | |
| try: | |
| with open(USAGE_FILE, "w") as f: | |
| json.dump(_usage, f) | |
| except Exception: | |
| pass | |
| def usage_today(): | |
| """Tokens used today (UTC); rolls the counter at the day boundary.""" | |
| if _usage["date"] != _today(): | |
| _usage["date"] = _today() | |
| _usage["used"] = 0 | |
| _save_usage() | |
| return _usage["used"] | |
| def usage_add(n): | |
| usage_today() | |
| _usage["used"] += max(0, int(n)) | |
| _save_usage() | |
| def estimate_prompt_tokens(body): | |
| n = 0 | |
| for m in body.get("messages", []): | |
| c = m.get("content") | |
| if isinstance(c, str): | |
| n += len(c) | |
| elif isinstance(c, list): | |
| n += sum(len(p.get("text", "")) for p in c if isinstance(p, dict)) | |
| return n // 4 + 8 | |
| _load_usage() | |
| def get_key(): | |
| key = os.environ.get("KIMI_API_KEY", "") | |
| if not key: | |
| raise ValueError("KIMI_API_KEY not set") | |
| return key | |
| def check_client_auth(request: Request): | |
| """Clients must present the proxy's own key. Returns a 401/503 response or None.""" | |
| proxy_key = os.environ.get("PROXY_KEY", "") | |
| if not proxy_key: | |
| return JSONResponse({"error": {"message": "PROXY_KEY not configured", "code": 503}}, status_code=503) | |
| auth = request.headers.get("Authorization", "") | |
| x_key = request.headers.get("x-api-key", "") | |
| if auth != f"Bearer {proxy_key}" and x_key != proxy_key: | |
| return JSONResponse({"error": {"message": "Invalid or missing API key", "code": 401}}, status_code=401) | |
| return None | |
| def auth_headers(): | |
| return {"Authorization": f"Bearer {get_key()}", "Content-Type": "application/json"} | |
| async def kimi_alive(): | |
| """Returns (ok, models_list | error_message).""" | |
| try: | |
| async with httpx.AsyncClient(timeout=8) as c: | |
| r = await c.get(f"{KIMI_BASE}/v1/models", headers=auth_headers()) | |
| if r.status_code == 200: | |
| return True, [m.get("id") for m in r.json().get("data", [])] | |
| return False, f"HTTP {r.status_code}: {r.text[:200]}" | |
| except Exception as e: | |
| return False, str(e) | |
| async def root(): | |
| if not os.environ.get("KIMI_API_KEY"): | |
| alive, info = False, "KIMI_API_KEY secret not set" | |
| else: | |
| alive, info = await kimi_alive() | |
| color = "#00ff88" if alive else "#ff4444" | |
| status = "ONLINE" if alive else "OFFLINE" | |
| models = ", ".join(info) if isinstance(info, list) else info | |
| used = usage_today() | |
| cap_color = "#ff4444" if used >= DAILY_CAP else "#00ff88" | |
| return HTMLResponse(f"""<html><body style="font-family:monospace;padding:2rem;background:#0d1117;color:#c9d1d9"> | |
| <h2>WirelessAudit Proxy</h2> | |
| <p>Backend: <code>Kimi K3</code> — <code>{KIMI_BASE}</code></p> | |
| <p>Status: <b style="color:{color}">{status}</b></p> | |
| <p>Models: <code>{models}</code></p> | |
| <p>Daily tokens (UTC): <b style="color:{cap_color}">{used:,}</b> / {DAILY_CAP:,} — resets 00:00 UTC</p> | |
| <p>Cline Base URL: <code>/v1</code> (requires API key)</p> | |
| </body></html>""") | |
| async def health(): | |
| try: | |
| get_key() | |
| except ValueError as e: | |
| return JSONResponse({"status": "offline", "error": str(e)}, status_code=503) | |
| alive, info = await kimi_alive() | |
| usage = {"used": usage_today(), "cap": DAILY_CAP} | |
| if alive: | |
| return {"status": "ok", "backend": "kimi", "models": info, "daily_tokens": usage} | |
| return JSONResponse({"status": "offline", "error": info, "daily_tokens": usage}, status_code=503) | |
| async def models(request: Request): | |
| if (deny := check_client_auth(request)) is not None: | |
| return deny | |
| try: | |
| async with httpx.AsyncClient(timeout=10) as c: | |
| r = await c.get(f"{KIMI_BASE}/v1/models", headers=auth_headers()) | |
| return JSONResponse(r.json(), status_code=r.status_code) | |
| except ValueError as e: | |
| return JSONResponse({"error": str(e)}, status_code=503) | |
| except Exception as e: | |
| return JSONResponse({"error": str(e)}, status_code=502) | |
| async def chat(request: Request): | |
| if (deny := check_client_auth(request)) is not None: | |
| return deny | |
| try: | |
| headers = auth_headers() | |
| except ValueError as e: | |
| return JSONResponse({"error": str(e)}, status_code=503) | |
| body = await request.json() | |
| if not body.get("model"): | |
| body["model"] = DEFAULT_MODEL | |
| stream = body.get("stream", False) | |
| est_prompt = estimate_prompt_tokens(body) | |
| used = usage_today() | |
| if used + est_prompt >= DAILY_CAP: | |
| return JSONResponse({"error": {"message": f"Daily token cap reached ({used:,}/{DAILY_CAP:,} tokens, UTC day). Resets at 00:00 UTC.", | |
| "type": "rate_limit_exceeded", "code": 429}}, status_code=429) | |
| url = f"{KIMI_BASE}/v1/chat/completions" | |
| timeout = httpx.Timeout(300.0, connect=15.0) | |
| try: | |
| if stream: | |
| so = dict(body.get("stream_options") or {}) | |
| so["include_usage"] = True | |
| body["stream_options"] = so | |
| c = httpx.AsyncClient(timeout=timeout) | |
| req = c.build_request("POST", url, json=body, headers=headers) | |
| r = await c.send(req, stream=True) | |
| if r.status_code != 200: | |
| err_body = await r.aread() | |
| code = r.status_code | |
| await r.aclose() | |
| await c.aclose() | |
| try: | |
| payload = json.loads(err_body) | |
| except Exception: | |
| payload = {"error": {"message": err_body.decode(errors="replace"), "code": code}} | |
| return JSONResponse(payload, status_code=code) | |
| async def proxy(): | |
| usage = None | |
| completion_chars = 0 | |
| buf = b"" | |
| try: | |
| async for chunk in r.aiter_bytes(): | |
| yield chunk | |
| buf += chunk | |
| if len(buf) > 65536: | |
| buf = buf[-4096:] | |
| while b"\n" in buf: | |
| line, buf = buf.split(b"\n", 1) | |
| line = line.strip() | |
| if not line.startswith(b"data:"): | |
| continue | |
| data = line[5:].strip() | |
| if not data or data == b"[DONE]": | |
| continue | |
| try: | |
| ev = json.loads(data) | |
| except Exception: | |
| continue | |
| if ev.get("usage"): | |
| usage = ev["usage"] | |
| for ch in ev.get("choices") or []: | |
| d = ch.get("delta") or {} | |
| completion_chars += len(d.get("content") or "") | |
| completion_chars += len(d.get("reasoning_content") or "") | |
| finally: | |
| await r.aclose() | |
| await c.aclose() | |
| if usage and usage.get("total_tokens"): | |
| usage_add(usage["total_tokens"]) | |
| else: | |
| usage_add(est_prompt + completion_chars // 4) | |
| return StreamingResponse(proxy(), media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) | |
| async with httpx.AsyncClient(timeout=timeout) as c: | |
| r = await c.post(url, json=body, headers=headers) | |
| try: | |
| payload = r.json() | |
| except Exception: | |
| return JSONResponse({"error": {"message": r.text, "code": r.status_code}}, | |
| status_code=r.status_code) | |
| if r.status_code == 200: | |
| u = payload.get("usage") or {} | |
| if u.get("total_tokens"): | |
| usage_add(u["total_tokens"]) | |
| else: | |
| out = "".join((ch.get("message") or {}).get("content") or "" | |
| for ch in payload.get("choices", [])) | |
| usage_add(est_prompt + len(out) // 4) | |
| return JSONResponse(payload, status_code=r.status_code) | |
| except httpx.ConnectError: | |
| return JSONResponse({"error": {"message": "Kimi unreachable", "code": 503}}, status_code=503) | |
| except httpx.ReadTimeout: | |
| return JSONResponse({"error": {"message": "Kimi timed out", "code": 504}}, status_code=504) | |
| except Exception as e: | |
| return JSONResponse({"error": {"message": str(e)}}, status_code=500) | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |