transformers - 💡(How to fix) Fix Improving CLI Serving Code Structure with Class-Based FastAPI Patterns [1 comments, 1 participants]

import uuid
from contextlib import asynccontextmanager

from ...utils import logging
from ...utils.import_utils import is_serve_available

if is_serve_available():
    from fastapi import FastAPI, Request, APIRouter, HTTPException
    from fastapi.middleware.cors import CORSMiddleware
    from fastapi.responses import JSONResponse, StreamingResponse
    from starlette.middleware.base import BaseHTTPMiddleware
    # For this example, I'm using a minimal library (`fastapi-cbx`) for the class-based routing.
    # The core architectural pattern of grouping endpoints in a class is independent of the implementation.
    from cbx import cbr

from .chat_completion import ChatCompletionHandler
from .completion import CompletionHandler
from .model_manager import ModelManager
from .response import ResponseHandler
from .transcription import TranscriptionHandler
from .utils import X_REQUEST_ID


@cbr(router=APIRouter())
class OpenAI:

    logger = logging.get_logger(__qualname__)

    def __init__(self, model_manager: ModelManager, chat_handler: ChatCompletionHandler, completion_handler: CompletionHandler, response_handler: ResponseHandler, transcription_handler: TranscriptionHandler):
        self.model_manager = model_manager
        self.chat_handler = chat_handler
        self.completion_handler = completion_handler
        self.response_handler = response_handler
        self.transcription_handler = transcription_handler

    @staticmethod
    def create_app(model_manager: ModelManager, chat_handler: ChatCompletionHandler, completion_handler: CompletionHandler, response_handler: ResponseHandler, transcription_handler: TranscriptionHandler, enable_cors: bool = False):
        service = OpenAI(
            model_manager,
            chat_handler,
            completion_handler,
            response_handler,
            transcription_handler
        )
        app = FastAPI(lifespan=service.lifespan)
        if enable_cors:
            app.add_middleware(
                CORSMiddleware,
                allow_origins=["*"],
                allow_credentials=True,
                allow_methods=["*"],
                allow_headers=["*"],
            )
            service.logger.warning_once(
                "CORS allow origin is set to `*`. Not recommended for production."
            )

        app.add_middleware(
            BaseHTTPMiddleware,
            dispatch=OpenAI.request_id_middleware
        )
        app.include_router(OpenAI.router)
        return app

    @asynccontextmanager
    async def lifespan(self, app: FastAPI):
        yield
        self.model_manager.shutdown()

    @staticmethod
    async def request_id_middleware(request: Request, call_next):
        """Get or set the request ID in the header."""
        request_id = request.headers.get(X_REQUEST_ID) or str(uuid.uuid4())
        request.state.request_id = request_id
        response = await call_next(request)
        response.headers[X_REQUEST_ID] = request_id
        return response

    @cbr.post("/v1/chat/completions")
    async def chat_completions(self, request: Request, body: dict):
        return await self.chat_handler.handle_request(body, request.state.request_id)

    @cbr.post("/v1/completions")
    async def completions(self, request: Request, body: dict):
        return await self.completion_handler.handle_request(body, request.state.request_id)

    @cbr.post("/v1/responses")
    async def responses(self, request: Request, body: dict):
        return await self.response_handler.handle_request(body, request.state.request_id)

    @cbr.post("/v1/audio/transcriptions")
    async def audio_transcriptions(self, request: Request):
        return await self.transcription_handler.handle_request(request)

    @cbr.post("/load_model")
    async def load_model(self, body: dict):
        model = body.get("model")
        if model is None:
            raise HTTPException(
                status_code=422,
                detail="Missing `model` field in the request body."
            )
        model_id_and_revision = self.model_manager.process_model_name(model)
        return StreamingResponse(
            self.model_manager.load_model_streaming(model_id_and_revision),
            media_type="text/event-stream"
        )

    @cbr.post("/reset")
    def reset(self):
        self.model_manager.shutdown()
        return JSONResponse({"status": "ok"})

    @cbr.get("/v1/models")
    @cbr.options("/v1/models")
    def list_models(self):
        return JSONResponse({"object": "list", "data": self.model_manager.get_gen_models()})

    @cbr.get("/health")
    @staticmethod
    def health():
        return JSONResponse({"status": "ok"})