BuckLakeAI / app.py
Parker's Fedora
Support anchored-path v2 serving.
c125dfd
Raw History Blame Contribute Delete
2.47 kB
from __future__ import annotations
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from bucklake_ai.config import get_settings
from bucklake_ai.errors import BuckLakeAIError, InvalidScalerArtifact, ModelArtifactNotFound, ScalerArtifactNotFound
from bucklake_ai.schemas import (
HealthResponse,
ModelMetaResponse,
PredictRequest,
PredictResponse,
)
from bucklake_ai.service import InferenceService
settings = get_settings()
service = InferenceService(settings)
@asynccontextmanager
async def lifespan(_: FastAPI):
logging.basicConfig(level=logging.INFO)
service.preload()
yield
app = FastAPI(title=settings.app_name, version=settings.app_version, lifespan=lifespan)
@app.exception_handler(ModelArtifactNotFound)
async def handle_model_not_found(_: Request, exc: ModelArtifactNotFound):
return JSONResponse(
status_code=503,
content={"error": "model_artifact_not_found", "detail": str(exc)},
)
@app.exception_handler(ScalerArtifactNotFound)
async def handle_scaler_not_found(_: Request, exc: ScalerArtifactNotFound):
return JSONResponse(
status_code=503,
content={"error": "scaler_artifact_not_found", "detail": str(exc)},
)
@app.exception_handler(InvalidScalerArtifact)
async def handle_invalid_scaler(_: Request, exc: InvalidScalerArtifact):
return JSONResponse(
status_code=503,
content={"error": "invalid_scaler_artifact", "detail": str(exc)},
)
@app.exception_handler(BuckLakeAIError)
async def handle_bucklake_error(_: Request, exc: BuckLakeAIError):
return JSONResponse(
status_code=400, content={"error": "bucklake_ai_error", "detail": str(exc)}
)
@app.exception_handler(Exception)
async def handle_unexpected_error(_: Request, exc: Exception):
return JSONResponse(
status_code=500, content={"error": "internal_error", "detail": str(exc)}
)
@app.get("/", response_model=ModelMetaResponse)
async def root():
return service.get_model_meta()
@app.get("/health", response_model=HealthResponse)
async def health():
return service.get_health()
@app.get("/api/model/meta", response_model=ModelMetaResponse)
async def model_meta():
return service.get_model_meta()
@app.post("/api/predict/v2", response_model=PredictResponse)
async def predict_v2(request: PredictRequest):
return service.predict(request)