File size: 4,176 Bytes
bb2b5cc | 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 | # Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
FastAPI app factory.
"""
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
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from .chat_completion import ChatCompletionHandler
from .model_manager import ModelManager
from .response import ResponseHandler
from .transcription import TranscriptionHandler
from .utils import X_REQUEST_ID
logger = logging.get_logger(__name__)
def build_server(
model_manager: ModelManager,
chat_handler: ChatCompletionHandler,
response_handler: ResponseHandler,
transcription_handler: TranscriptionHandler,
enable_cors: bool = False,
) -> FastAPI:
"""Build and return a configured FastAPI application.
Args:
model_manager: Handles model loading, caching, and cleanup.
chat_handler: Handles `/v1/chat/completions` requests.
response_handler: Handles `/v1/responses` requests.
enable_cors: If `True`, adds permissive CORS middleware (allow all origins).
Returns:
A FastAPI app ready to be passed to uvicorn.
"""
@asynccontextmanager
async def lifespan(app: FastAPI):
yield
model_manager.shutdown()
app = FastAPI(lifespan=lifespan)
if enable_cors:
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
logger.warning_once("CORS allow origin is set to `*`. Not recommended for production.")
# ---- Middleware ----
@app.middleware("http")
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
# ---- Routes ----
@app.post("/v1/chat/completions")
async def chat_completions(request: Request, body: dict):
return await chat_handler.handle_request(body, request.state.request_id)
@app.post("/v1/responses")
async def responses(request: Request, body: dict):
return await response_handler.handle_request(body, request.state.request_id)
@app.post("/v1/audio/transcriptions")
async def audio_transcriptions(request: Request):
return await transcription_handler.handle_request(request)
@app.post("/load_model")
async def load_model(body: dict):
from fastapi import HTTPException
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 = model_manager.process_model_name(model)
return StreamingResponse(
model_manager.load_model_streaming(model_id_and_revision), media_type="text/event-stream"
)
@app.post("/reset")
def reset():
model_manager.shutdown()
return JSONResponse({"status": "ok"})
@app.get("/v1/models")
@app.options("/v1/models")
def list_models():
return JSONResponse({"object": "list", "data": model_manager.get_gen_models()})
@app.get("/health")
def health():
return JSONResponse({"status": "ok"})
return app
|