Model/api/developer_gateway.py

"""
Frox AI — Developer API Gateway

A thin, OpenAI-compatible gateway that sits IN FRONT of the Morph model
server (api/server.py) and adds the pieces a public developer API needs
but the raw model server intentionally doesn't have:

  * API key issuance + revocation      (POST/DELETE /dev/keys)
  * Bearer-token authentication        (Authorization: Bearer frx-...)
  * Per-key rate limiting              (requests/min, token budgets)
  * Usage tracking                     (GET /dev/usage)

It forwards authenticated /v1/* calls to the upstream Morph server
unchanged — so every existing OpenAI-compatible client keeps working,
they just point at the gateway's URL and send their Frox API key.

Run (gateway on 8080, model server on 8000):
    python scripts/serve.py --family classic --model ./ckpt --port 8000
    MORPH_UPSTREAM_URL=http://localhost:8000 \
        uvicorn api.developer_gateway:app --host 0.0.0.0 --port 8080

Storage is a local JSON file by default (MORPH_KEYS_PATH). Swap
`KeyStore` for a real database in production — the interface is small.
"""
from __future__ import annotations

import json
import os
import secrets
import threading
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Deque, Dict, Optional

import httpx
from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field


UPSTREAM_URL = os.environ.get("MORPH_UPSTREAM_URL", "http://localhost:8000").rstrip("/")
KEYS_PATH = os.environ.get("MORPH_KEYS_PATH", "./data/api_keys.json")
KEY_PREFIX = "frx-"


# ── Key model + storage ─────────────────────────────────────────────

@dataclass
class ApiKey:
    key: str
    name: str
    plan: str = "free"
    created_at: float = field(default_factory=time.time)
    revoked: bool = False
    # usage counters (cumulative)
    total_requests: int = 0
    total_prompt_tokens: int = 0
    total_completion_tokens: int = 0

    # per-plan limits
    @property
    def rpm_limit(self) -> int:
        return {"free": 20, "pro": 120, "scale": 600}.get(self.plan, 20)


class KeyStore:
    """JSON-backed key store. Thread-safe for the gateway's needs."""

    def __init__(self, path: str):
        self.path = Path(path)
        self._lock = threading.Lock()
        self._keys: Dict[str, ApiKey] = {}
        self._load()

    def _load(self):
        if self.path.exists():
            try:
                raw = json.loads(self.path.read_text())
                self._keys = {k: ApiKey(**v) for k, v in raw.items()}
            except (json.JSONDecodeError, TypeError):
                self._keys = {}

    def _flush(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        tmp.write_text(json.dumps({k: asdict(v) for k, v in self._keys.items()}, indent=2))
        tmp.replace(self.path)

    def create(self, name: str, plan: str = "free") -> ApiKey:
        with self._lock:
            token = KEY_PREFIX + secrets.token_urlsafe(32)
            api_key = ApiKey(key=token, name=name, plan=plan)
            self._keys[token] = api_key
            self._flush()
            return api_key

    def get(self, token: str) -> Optional[ApiKey]:
        return self._keys.get(token)

    def revoke(self, token: str) -> bool:
        with self._lock:
            k = self._keys.get(token)
            if not k:
                return False
            k.revoked = True
            self._flush()
            return True

    def record_usage(self, token: str, prompt_tokens: int, completion_tokens: int):
        with self._lock:
            k = self._keys.get(token)
            if not k:
                return
            k.total_requests += 1
            k.total_prompt_tokens += prompt_tokens
            k.total_completion_tokens += completion_tokens
            self._flush()

    def all(self):
        return list(self._keys.values())


store = KeyStore(KEYS_PATH)


# ── Rate limiter (sliding window, in-memory) ────────────────────────────

class RateLimiter:
    def __init__(self):
        self._hits: Dict[str, Deque[float]] = defaultdict(deque)
        self._lock = threading.Lock()

    def check(self, token: str, limit_per_min: int) -> bool:
        now = time.time()
        window_start = now - 60
        with self._lock:
            hits = self._hits[token]
            while hits and hits[0] < window_start:
                hits.popleft()
            if len(hits) >= limit_per_min:
                return False
            hits.append(now)
            return True


limiter = RateLimiter()


# ── Auth dependency ──────────────────────────────────────────────

def _admin_ok(x_admin_token: Optional[str]) -> bool:
    expected = os.environ.get("MORPH_ADMIN_TOKEN")
    # If no admin token is configured, key management is open on localhost
    # only — fine for local dev, must be set in production.
    return expected is None or x_admin_token == expected


async def require_api_key(authorization: Optional[str] = Header(None)) -> ApiKey:
    if not authorization or not authorization.lower().startswith("bearer "):
        raise HTTPException(status_code=401, detail={
            "error": {"message": "Missing bearer token. Send 'Authorization: Bearer frx-...'",
                      "type": "unauthorized"}})
    token = authorization[7:].strip()
    api_key = store.get(token)
    if not api_key or api_key.revoked:
        raise HTTPException(status_code=401, detail={
            "error": {"message": "Invalid or revoked API key.", "type": "unauthorized"}})
    if not limiter.check(token, api_key.rpm_limit):
        raise HTTPException(status_code=429, detail={
            "error": {"message": f"Rate limit exceeded ({api_key.rpm_limit} req/min for plan "
                                 f"'{api_key.plan}').", "type": "rate_limit_exceeded"}})
    return api_key


# ── App ────────────────────────────────────────────────────

app = FastAPI(title="Frox AI — Developer API Gateway", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.environ.get("MORPH_CORS_ORIGINS", "*").split(","),
    allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
)

_client: Optional[httpx.AsyncClient] = None


@app.on_event("startup")
async def _startup():
    global _client
    _client = httpx.AsyncClient(base_url=UPSTREAM_URL, timeout=httpx.Timeout(300.0))


@app.on_event("shutdown")
async def _shutdown():
    if _client:
        await _client.aclose()


# ── Key management (admin) ───────────────────────────────────────

class CreateKeyRequest(BaseModel):
    name: str = Field(..., description="Human label for the key")
    plan: str = Field("free", description="free | pro | scale")


@app.post("/dev/keys")
async def create_key(body: CreateKeyRequest, x_admin_token: Optional[str] = Header(None)):
    if not _admin_ok(x_admin_token):
        raise HTTPException(status_code=403, detail={"error": {"message": "Admin token required"}})
    key = store.create(body.name, body.plan)
    # The full secret is returned exactly once, at creation time.
    return {"key": key.key, "name": key.name, "plan": key.plan, "rpm_limit": key.rpm_limit,
            "created_at": key.created_at}


@app.delete("/dev/keys/{token}")
async def revoke_key(token: str, x_admin_token: Optional[str] = Header(None)):
    if not _admin_ok(x_admin_token):
        raise HTTPException(status_code=403, detail={"error": {"message": "Admin token required"}})
    if not store.revoke(token):
        raise HTTPException(status_code=404, detail={"error": {"message": "Key not found"}})
    return {"revoked": True, "key": token}


@app.get("/dev/usage")
async def usage(api_key: ApiKey = Depends(require_api_key)):
    return {
        "name": api_key.name, "plan": api_key.plan, "rpm_limit": api_key.rpm_limit,
        "total_requests": api_key.total_requests,
        "total_prompt_tokens": api_key.total_prompt_tokens,
        "total_completion_tokens": api_key.total_completion_tokens,
    }


# ── OpenAI-compatible proxy (authenticated) ────────────────────────────

@app.get("/v1/models")
async def proxy_models(api_key: ApiKey = Depends(require_api_key)):
    resp = await _client.get("/v1/models")
    return JSONResponse(status_code=resp.status_code, content=resp.json())


@app.post("/v1/chat/completions")
async def proxy_chat(request: Request, api_key: ApiKey = Depends(require_api_key)):
    body = await request.body()
    try:
        payload = json.loads(body or b"{}")
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail={"error": {"message": "Invalid JSON body"}})

    is_stream = bool(payload.get("stream"))

    if is_stream:
        # Pass the SSE stream straight through; usage is recorded as one
        # request (token accounting for streamed responses would require
        # parsing the final usage chunk — counted as a request here).
        async def _relay():
            async with _client.stream("POST", "/v1/chat/completions", content=body,
                                      headers={"Content-Type": "application/json"}) as upstream:
                async for chunk in upstream.aiter_raw():
                    yield chunk
        store.record_usage(api_key.key, 0, 0)
        return StreamingResponse(_relay(), media_type="text/event-stream")

    resp = await _client.post("/v1/chat/completions", content=body,
                              headers={"Content-Type": "application/json"})
    data = resp.json()
    u = data.get("usage") or {}
    store.record_usage(api_key.key, u.get("prompt_tokens", 0), u.get("completion_tokens", 0))
    return JSONResponse(status_code=resp.status_code, content=data)


@app.get("/health")
async def health():
    try:
        r = await _client.get("/health")
        return {"gateway": "ok", "upstream": r.json()}
    except Exception as e:
        return JSONResponse(status_code=503, content={"gateway": "ok", "upstream": f"unreachable: {e}"})