server.py

#!/usr/bin/env python3
"""HTTP API around Emma-5 — host-agnostic (local, tunnel, HF Space, Cloud Run).

Run locally:
    ./.venv/bin/uvicorn server:app --port 8000
Then POST /chat {"question": "..."} -> {"answer": "...", "seconds": 3.1}
"""
import hmac
import os
import threading
import time

from fastapi import Depends, FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from emma_chat import Emma, DEFAULT_MODEL_DIR

# Optional shared-secret auth. When EMMA_API_KEY is set (e.g. on a public VM),
# /chat requires a matching X-API-Key header. Unset (local dev) => open.
API_KEY = os.environ.get("EMMA_API_KEY")


def require_key(x_api_key: str | None = Header(default=None)):
    if API_KEY and not (x_api_key and hmac.compare_digest(x_api_key, API_KEY)):
        raise HTTPException(status_code=401, detail="invalid or missing api key")

app = FastAPI(title="Emma4ever — live Emma-5")

app.add_middleware(
    CORSMiddleware,
    allow_origins=os.environ.get("EMMA_ALLOWED_ORIGINS", "*").split(","),
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load once at import (HF Spaces / uvicorn workers each load their own copy).
_emma = Emma(os.environ.get("EMMA_MODEL_DIR", DEFAULT_MODEL_DIR))
# Serialize inference: one CPU forward at a time keeps memory + latency bounded
# on small free instances.
_lock = threading.Lock()


class ChatRequest(BaseModel):
    question: str = Field(min_length=1, max_length=500)
    temperature: float = Field(0.85, ge=0.0, le=2.0)
    top_p: float = Field(0.95, ge=0.1, le=1.0)
    max_new: int = Field(80, ge=1, le=200)
    greedy: bool = False


class ChatResponse(BaseModel):
    answer: str
    seconds: float


@app.get("/health")
def health():
    return {"ok": True, "model": "emma-5"}


@app.post("/chat", response_model=ChatResponse, dependencies=[Depends(require_key)])
def chat(req: ChatRequest):
    t0 = time.time()
    with _lock:
        answer = _emma.generate(
            req.question,
            max_new=req.max_new,
            temperature=req.temperature,
            top_p=req.top_p,
            greedy=req.greedy,
        )
    return ChatResponse(answer=answer, seconds=round(time.time() - t0, 2))