| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import importlib |
| 4 | import inspect |
| 5 | from argparse import Namespace |
| 6 | |
| 7 | from fastapi import FastAPI |
| 8 | from starlette.middleware.cors import CORSMiddleware |
| 9 | |
| 10 | from vllm import envs |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.tasks import SupportedTask |
| 13 | |
| 14 | from .log_response import log_response |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | def init_entrypoints_middleware( |
| 20 | args: Namespace, |
| 21 | app: FastAPI, |
| 22 | supported_tasks: tuple["SupportedTask", ...], |
| 23 | ): |
| 24 | app.add_middleware( |
| 25 | CORSMiddleware, |
| 26 | allow_origins=args.allowed_origins, |
| 27 | allow_credentials=args.allow_credentials, |
| 28 | allow_methods=args.allowed_methods, |
| 29 | allow_headers=args.allowed_headers, |
| 30 | ) |
| 31 | |
| 32 | # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY |
| 33 | if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]: |
| 34 | from .authenticate import AuthenticationMiddleware |
| 35 | |
| 36 | app.add_middleware(AuthenticationMiddleware, tokens=tokens) |
| 37 | |
| 38 | if args.enable_request_id_headers: |
| 39 | from .x_request_id import XRequestIdMiddleware |
| 40 | |
| 41 | app.add_middleware(XRequestIdMiddleware) |
| 42 | |
| 43 | if "generate" in supported_tasks: |
| 44 | # Add scaling middleware to check for scaling state |
| 45 | from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware |
| 46 | |
| 47 | app.add_middleware(ScalingMiddleware) |
| 48 | |
| 49 | if "realtime" in supported_tasks: |
| 50 | # Add WebSocket metrics middleware |
| 51 | from vllm.entrypoints.speech_to_text.realtime.metrics import ( |
| 52 | WebSocketMetricsMiddleware, |
| 53 | ) |
| 54 | |
| 55 | app.add_middleware(WebSocketMetricsMiddleware) |
| 56 | |
| 57 | if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE: |
| 58 | logger.warning( |
| 59 | "CAUTION: Enabling log response in the API Server. " |
| 60 | "This can include sensitive information and should be " |
| 61 | "avoided in production." |
| 62 | ) |
| 63 | app.middleware("http")(log_response) |
| 64 | |
| 65 | for middleware in args.middleware: |
| 66 | module_path, object_name = middleware.rsplit(".", 1) |
| 67 | imported = getattr(importlib.import_module(module_path), object_name) |
| 68 | if inspect.isclass(imported): |
| 69 | app.add_middleware(imported) # type: ignore[arg-type] |
| 70 | elif inspect.iscoroutinefunction(imported): |
| 71 | app.middleware("http")(imported) |
| 72 | else: |
| 73 | raise ValueError( |
| 74 | f"Invalid middleware {middleware}. Must be a function or a class." |
| 75 | ) |