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)
@app.get("/")
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"""
WirelessAudit Proxy
Backend: Kimi K3 — {KIMI_BASE}
Status: {status}
Models: {models}
Daily tokens (UTC): {used:,} / {DAILY_CAP:,} — resets 00:00 UTC
Cline Base URL: /v1 (requires API key)
""")
@app.get("/health")
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)
@app.get("/v1/models")
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)
@app.post("/v1/chat/completions")
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)