File size: 15,356 Bytes
0cd599e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 | 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:
# Joining a slow writer must not block the event loop.
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):
# 无内置权重镜像：模型在子目录 laya-<版本>-<精度>/ 下
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
# uvicorn 经 server.main:app 引用；pytest 导入本模块时跳过（测试自建 Settings + create_app）
if "pytest" not in sys.modules:
app = create_app()
|