| import hashlib
|
| import hmac
|
| import logging
|
| import secrets
|
| import sys
|
| import threading
|
| import time
|
| from contextlib import asynccontextmanager
|
| from datetime import datetime, timedelta, timezone
|
| from typing import Annotated, Any
|
|
|
| from argon2 import PasswordHasher
|
| from argon2.exceptions import VerificationError
|
| from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response
|
| from fastapi.exceptions import RequestValidationError
|
| from fastapi.responses import JSONResponse
|
| from starlette.concurrency import run_in_threadpool
|
|
|
| from .ad_classifier import (AdClassifierPredictor, MultiPrecisionPredictor,
|
| Predictor, served_model_ids)
|
| from .config import Settings
|
| from .database import Database, utc_now
|
| from .model_manager import ModelManager
|
| from .schemas import (CreateKeyRequest, InferenceRequest, LoginRequest,
|
| ModelDownloadRequest)
|
| from .usage_writer import UsageEvent, UsageWriter
|
|
|
| logger = logging.getLogger(__name__)
|
| SESSION_COOKIE = "laya_session"
|
| CSRF_COOKIE = "laya_csrf"
|
|
|
|
|
| def digest(value: str) -> str:
|
| return hashlib.sha256(value.encode()).hexdigest()
|
|
|
|
|
| def error(status: int, code: str, message: str) -> HTTPException:
|
| return HTTPException(status_code=status, detail={"code": code, "message": message})
|
|
|
|
|
| def create_app(settings: Settings | None = None, predictor: Predictor | None = None) -> FastAPI:
|
| settings = settings or Settings.from_env()
|
| db = Database(settings.database_path)
|
| db.initialize()
|
| usage_writer = UsageWriter(db)
|
|
|
| @asynccontextmanager
|
| async def lifespan(app: FastAPI):
|
| usage_writer.start()
|
| try:
|
| yield
|
| finally:
|
|
|
| await run_in_threadpool(usage_writer.stop)
|
|
|
| app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None, lifespan=lifespan)
|
| hasher = PasswordHasher()
|
| router_holder: dict[str, Predictor] = {}
|
| router_lock = threading.Lock()
|
| manager = None if (settings.model_dir / "MODEL_ID.txt").is_file() \
|
| else ModelManager(settings.model_dir, settings.device)
|
|
|
| @app.exception_handler(RequestValidationError)
|
| def validation_error(request: Request, exc: RequestValidationError) -> JSONResponse:
|
| fields = [{"loc": list(item["loc"]), "type": item["type"], "message": item["msg"]}
|
| for item in exc.errors()]
|
| return JSONResponse(status_code=422, content={"detail": {
|
| "code": "VALIDATION_ERROR", "message": "Request validation failed", "errors": fields,
|
| }})
|
|
|
| def current_session(request: Request) -> str:
|
| token = request.cookies.get(SESSION_COOKIE)
|
| if not token:
|
| raise error(401, "UNAUTHENTICATED", "Sign in required")
|
| with db.connect() as connection:
|
| row = connection.execute(
|
| "SELECT digest FROM sessions WHERE digest=? AND expires_at>?",
|
| (digest(token), utc_now()),
|
| ).fetchone()
|
| if row is None:
|
| raise error(401, "UNAUTHENTICATED", "Sign in required")
|
| return token
|
|
|
| Session = Annotated[str, Depends(current_session)]
|
|
|
| def require_csrf(request: Request, session: Session, x_csrf_token: Annotated[str | None, Header()] = None) -> None:
|
| cookie = request.cookies.get(CSRF_COOKIE)
|
| if not cookie or not x_csrf_token or not hmac.compare_digest(cookie, x_csrf_token):
|
| raise error(403, "CSRF_INVALID", "CSRF token is invalid")
|
| with db.connect() as connection:
|
| row = connection.execute("SELECT csrf_digest FROM sessions WHERE digest=?", (digest(session),)).fetchone()
|
| if row is None or not hmac.compare_digest(row["csrf_digest"], digest(cookie)):
|
| raise error(403, "CSRF_INVALID", "CSRF token is invalid")
|
|
|
| WriteSession = Annotated[None, Depends(require_csrf)]
|
|
|
| def identify_key(authorization: str | None) -> int:
|
| if not authorization or not authorization.startswith("Bearer "):
|
| raise error(401, "INVALID_API_KEY", "Valid Bearer API key required")
|
| token = authorization[7:]
|
| if not token.startswith("laya_") or len(token) < 32:
|
| raise error(401, "INVALID_API_KEY", "Valid Bearer API key required")
|
| with db.connect() as connection:
|
| row = connection.execute(
|
| "SELECT id FROM api_keys WHERE digest=? AND revoked_at IS NULL", (digest(token),)
|
| ).fetchone()
|
| if row is None:
|
| raise error(401, "INVALID_API_KEY", "Valid Bearer API key required")
|
| return row["id"]
|
|
|
| def api_key_id(authorization: Annotated[str | None, Header()] = None) -> int:
|
| return identify_key(authorization)
|
|
|
| ApiKeyId = Annotated[int, Depends(api_key_id)]
|
|
|
| def infer(payload: InferenceRequest, source: str, key_id: int | None) -> dict[str, Any]:
|
| with router_lock:
|
| if "router" not in router_holder:
|
| try:
|
| if (settings.model_dir / "MODEL_ID.txt").is_file():
|
| router_holder["router"] = predictor or AdClassifierPredictor(
|
| settings.model_dir, settings.device
|
| )
|
| else:
|
| router_holder["router"] = predictor or MultiPrecisionPredictor(
|
| settings.model_dir, settings.device
|
| )
|
| except Exception:
|
| logger.exception("Failed to initialize ad-clf predictor")
|
| raise error(503, "MODEL_UNAVAILABLE", "Model is unavailable")
|
| served = router_holder["router"].served_ids()
|
| if payload.model not in served:
|
| raise error(422, "MODEL_NOT_AVAILABLE", "Installed models: " + ", ".join(served))
|
| started = time.monotonic()
|
| try:
|
| overrides = {} if payload.max_len is None else {"max_len": payload.max_len}
|
| result = router_holder["router"].predict(
|
| payload.state,
|
| {qid: q.model_dump(exclude_none=True) for qid, q in payload.questions.items()},
|
| model=None if payload.model == "auto" else payload.model,
|
| **overrides,
|
| )
|
| if not isinstance(result, dict) or not isinstance(result.get("answers"), dict):
|
| raise ValueError("Laya returned invalid answers")
|
| usage = result["usage"]
|
| if not isinstance(usage["input_tokens"], int) or not isinstance(usage["output_tokens"], int):
|
| raise ValueError("Laya returned invalid usage")
|
| if usage["input_tokens"] < 0 or usage["output_tokens"] < 0:
|
| raise ValueError("Laya returned negative usage")
|
| if not isinstance(result.get("model"), str):
|
| raise ValueError("Laya returned invalid model")
|
| except Exception as exc:
|
| logger.error("Laya inference failed: %s", type(exc).__name__)
|
| raise error(503, "MODEL_UNAVAILABLE", "Model is unavailable")
|
| duration_ms = round((time.monotonic() - started) * 1000)
|
| usage_writer.submit(UsageEvent(
|
| utc_now(), source, key_id, result["model"],
|
| usage["input_tokens"], usage["output_tokens"], duration_ms,
|
| ))
|
| return result
|
|
|
| @app.get("/health/live")
|
| def live() -> dict[str, str]:
|
| return {"status": "ok"}
|
|
|
| @app.get("/health/ready")
|
| def ready() -> dict[str, str]:
|
| def has_model(d) -> bool:
|
| return ((d / "model.safetensors").is_file()
|
| and (d / "config.json").is_file()
|
| and (d / "tokenizer.json").is_file())
|
|
|
| root = settings.model_dir
|
| if not has_model(root):
|
|
|
| found = root.is_dir() and any(has_model(d) for d in root.iterdir() if d.is_dir())
|
| if not found:
|
| raise error(503, "MODEL_UNAVAILABLE", "ad-clf model files are missing")
|
| return {"status": "ready"}
|
|
|
| @app.post("/internal/auth/login")
|
| def login(payload: LoginRequest, request: Request, response: Response) -> dict[str, str]:
|
| source = request.client.host if request.client else "unknown"
|
| since = (datetime.now(timezone.utc) - timedelta(minutes=15)).isoformat(timespec="seconds")
|
| with db.connect() as connection:
|
| connection.execute("DELETE FROM login_attempts WHERE attempted_at<?", (since,))
|
| attempts = connection.execute(
|
| "SELECT COUNT(*) FROM login_attempts WHERE source=? AND username=?",
|
| (source, payload.username),
|
| ).fetchone()[0]
|
| if attempts >= 5:
|
| raise error(429, "LOGIN_RATE_LIMITED", "Too many login attempts. Try again later")
|
| connection.execute("INSERT INTO login_attempts VALUES (?,?,?)", (source, payload.username, utc_now()))
|
| try:
|
| valid_password = hasher.verify(settings.admin_password_hash, payload.password)
|
| except VerificationError:
|
| valid_password = False
|
| if payload.username != settings.admin_username or not valid_password:
|
| raise error(401, "INVALID_CREDENTIALS", "Invalid username or password")
|
| with db.connect() as connection:
|
| connection.execute("DELETE FROM login_attempts WHERE source=? AND username=?", (source, payload.username))
|
| session = secrets.token_urlsafe(48)
|
| csrf = secrets.token_urlsafe(32)
|
| expires = (datetime.now(timezone.utc) + timedelta(hours=settings.session_hours)).isoformat(timespec="seconds")
|
| with db.connect() as connection:
|
| connection.execute("INSERT INTO sessions VALUES (?,?,?)", (digest(session), digest(csrf), expires))
|
| is_https = request.url.scheme == "https" or request.headers.get("origin", "").startswith("https://")
|
| cookie_options = {"secure": is_https, "samesite": "lax", "path": "/", "max_age": settings.session_hours * 3600}
|
| response.set_cookie(SESSION_COOKIE, session, httponly=True, **cookie_options)
|
| response.set_cookie(CSRF_COOKIE, csrf, httponly=False, **cookie_options)
|
| return {"username": settings.admin_username}
|
|
|
| @app.get("/internal/auth/session")
|
| def session(request: Request, session: Session) -> dict[str, str]:
|
| return {"username": settings.admin_username, "csrf_token": request.cookies.get(CSRF_COOKIE, "")}
|
|
|
| @app.post("/internal/auth/logout")
|
| def logout(request: Request, response: Response, session: Session, _: WriteSession) -> dict[str, bool]:
|
| with db.connect() as connection:
|
| connection.execute("DELETE FROM sessions WHERE digest=?", (digest(session),))
|
| response.delete_cookie(SESSION_COOKIE, path="/")
|
| response.delete_cookie(CSRF_COOKIE, path="/")
|
| return {"ok": True}
|
|
|
| @app.get("/internal/api-keys")
|
| def list_keys(session: Session) -> list[dict[str, Any]]:
|
| with db.connect() as connection:
|
| rows = connection.execute(
|
| "SELECT id,name,mask,created_at,revoked_at,last_used_at FROM api_keys ORDER BY id DESC"
|
| ).fetchall()
|
| return [dict(row) for row in rows]
|
|
|
| @app.get("/internal/models")
|
| def available_models(session: Session) -> dict[str, list[str]]:
|
| router = router_holder.get("router")
|
| if router is not None and hasattr(router, "served_ids"):
|
| return {"models": router.served_ids()}
|
| return {"models": served_model_ids(settings.model_dir)}
|
|
|
| @app.get("/internal/models/catalog")
|
| def model_catalog(session: Session) -> dict[str, Any]:
|
| if manager is None:
|
| raise error(422, "MODEL_NOT_DOWNLOADABLE", "Baked image: model is built in")
|
| return manager.catalog()
|
|
|
| @app.post("/internal/models/download")
|
| def download_model(payload: ModelDownloadRequest, session: Session, _: WriteSession) -> dict[str, Any]:
|
| if manager is None:
|
| raise error(422, "MODEL_NOT_DOWNLOADABLE", "Baked image: model is built in")
|
| return manager.download(payload.precision)
|
|
|
| @app.get("/internal/models/download/status")
|
| def download_status(session: Session) -> dict[str, Any]:
|
| if manager is None:
|
| return {"state": "baked", "precision": None, "error": None}
|
| return manager.status()
|
|
|
| @app.post("/internal/api-keys", status_code=201)
|
| def create_key(payload: CreateKeyRequest, session: Session, _: WriteSession) -> dict[str, Any]:
|
| token = "laya_" + secrets.token_urlsafe(32)
|
| mask = "laya_••••" + token[-6:]
|
| with db.connect() as connection:
|
| cursor = connection.execute(
|
| "INSERT INTO api_keys (name,digest,mask,created_at) VALUES (?,?,?,?)",
|
| (payload.name, digest(token), mask, utc_now()),
|
| )
|
| return {"id": cursor.lastrowid, "name": payload.name, "key": token, "mask": mask}
|
|
|
| @app.post("/internal/api-keys/{key_id}/revoke")
|
| def revoke_key(key_id: int, session: Session, _: WriteSession) -> dict[str, bool]:
|
| with db.connect() as connection:
|
| cursor = connection.execute(
|
| "UPDATE api_keys SET revoked_at=? WHERE id=? AND revoked_at IS NULL",
|
| (utc_now(), key_id),
|
| )
|
| if cursor.rowcount == 0:
|
| raise error(404, "KEY_NOT_FOUND", "Active API key not found")
|
| return {"ok": True}
|
|
|
| @app.post("/v1/systemone")
|
| def systemone(key_id: ApiKeyId, payload: InferenceRequest) -> dict[str, Any]:
|
| return infer(payload, "api_key", key_id)
|
|
|
| @app.post("/internal/playground")
|
| def playground(payload: InferenceRequest, session: Session, _: WriteSession) -> dict[str, Any]:
|
| return infer(payload, "playground", None)
|
|
|
| @app.get("/internal/usage")
|
| def usage(session: Session) -> dict[str, Any]:
|
| with db.connect() as connection:
|
| totals = connection.execute(
|
| "SELECT COUNT(*) requests, COALESCE(SUM(input_tokens),0) input_tokens, COALESCE(SUM(output_tokens),0) output_tokens FROM usage_events"
|
| ).fetchone()
|
| daily = connection.execute(
|
| "SELECT substr(created_at,1,10) day, COUNT(*) requests, SUM(input_tokens) input_tokens, SUM(output_tokens) output_tokens FROM usage_events GROUP BY day ORDER BY day"
|
| ).fetchall()
|
| sources = connection.execute(
|
| "SELECT source, key_id, COUNT(*) requests, SUM(input_tokens) input_tokens, SUM(output_tokens) output_tokens FROM usage_events GROUP BY source,key_id ORDER BY requests DESC"
|
| ).fetchall()
|
| return {"totals": dict(totals), "daily": [dict(row) for row in daily], "sources": [dict(row) for row in sources]}
|
|
|
| if settings.frontend_dir.is_dir():
|
| app.frontend("/", directory=str(settings.frontend_dir))
|
| return app
|
|
|
|
|
|
|
| if "pytest" not in sys.modules:
|
| app = create_app()
|
|
|
