diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py new file mode 100644 index 0000000000..ae2481384b --- /dev/null +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -0,0 +1,39 @@ +""" +Prometheus Auth Middleware +""" +from fastapi import Request +from starlette.middleware.base import BaseHTTPMiddleware + +import litellm +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + +class PrometheusAuthMiddleware(BaseHTTPMiddleware): + async def dispatch(self, request: Request, call_next): + # Check if this is a request to the metrics endpoint + + if self._is_prometheus_metrics_endpoint(request): + try: + await user_api_key_auth( + request=request, api_key=request.headers.get("Authorization") or "" + ) + except Exception as e: + raise e + + # Process the request and get the response + response = await call_next(request) + + return response + + @staticmethod + def _is_prometheus_metrics_endpoint(request: Request): + try: + if "/metrics" in request.url.path: + return True + return False + except Exception: + return False + + @staticmethod + def _should_run_auth_on_metrics_endpoint(): + return litellm.require_auth_for_metrics_endpoint diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1d2eab31c4..32fa4e8f52 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -249,6 +249,7 @@ from litellm.proxy.management_endpoints.ui_sso import ( ) from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) @@ -745,6 +746,7 @@ app.add_middleware( allow_headers=["*"], ) +app.add_middleware(PrometheusAuthMiddleware) from typing import Dict