From 253060cb092534baeaaf2b5d130e318cc1d4eb65 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 17:35:02 -0700 Subject: [PATCH 1/7] allow requiring auth for /metrics endpoint --- litellm/__init__.py | 1 + litellm/integrations/prometheus.py | 74 ++++++++++++++++++++ litellm/proxy/common_utils/callback_utils.py | 13 +--- litellm/proxy/proxy_config.yaml | 14 ++-- litellm/proxy/proxy_server.py | 9 +-- 5 files changed, 87 insertions(+), 24 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 42a96abf134..c9d9f3aaf2f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -123,6 +123,7 @@ callbacks: List[ langfuse_default_tags: Optional[List[str]] = None langsmith_batch_size: Optional[int] = None prometheus_initialize_budget_metrics: Optional[bool] = False +require_auth_for_metrics_endpoint: Optional[bool] = False argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload gcs_pub_sub_use_v1: Optional[ diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 5ac8c80eb30..6cd413857bf 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1721,6 +1721,80 @@ class PrometheusLogger(CustomLogger): return (end_time - start_time).total_seconds() return None + @staticmethod + def _mount_metrics_endpoint(premium_user: bool): + """ + Mount the Prometheus metrics endpoint with optional authentication. + + Args: + premium_user (bool): Whether the user is a premium user + require_auth (bool, optional): Whether to require authentication for the metrics endpoint. + Defaults to False. + """ + from prometheus_client import make_asgi_app + + from litellm._logging import verbose_proxy_logger + from litellm.proxy._types import CommonProxyErrors + from litellm.proxy.proxy_server import app + + if premium_user is not True: + verbose_proxy_logger.warning( + f"Prometheus metrics are only available for premium users. {CommonProxyErrors.not_premium_user.value}" + ) + + if PrometheusLogger._should_init_metrics_with_auth(): + PrometheusLogger._mount_metrics_endpoint_with_auth() + else: + # Mount metrics directly without authentication + PrometheusLogger._mount_metrics_endpoint_without_auth() + + @staticmethod + def _mount_metrics_endpoint_without_auth(): + from prometheus_client import make_asgi_app + + from litellm._logging import verbose_proxy_logger + from litellm.proxy.proxy_server import app + + # Create metrics ASGI app + metrics_app = make_asgi_app() + + # Mount the metrics app to the app + app.mount("/metrics", metrics_app) + verbose_proxy_logger.debug( + "Starting Prometheus Metrics on /metrics (no authentication)" + ) + + @staticmethod + def _mount_metrics_endpoint_with_auth(): + from fastapi import APIRouter, Depends + from prometheus_client import make_asgi_app + + from litellm._logging import verbose_proxy_logger + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + # Create metrics ASGI app + metrics_app = make_asgi_app() + + # Create a router for authenticated metrics + metrics_router = APIRouter() + + # Add metrics endpoint with authentication + @metrics_router.get("/metrics", dependencies=[Depends(user_api_key_auth)]) + async def authenticated_metrics(): + verbose_proxy_logger.debug("Serving authenticated metrics endpoint") + return metrics_app + + # Mount the router to the app + app.include_router(metrics_router) + verbose_proxy_logger.debug( + "Starting Prometheus Metrics on /metrics with authentication" + ) + + @staticmethod + def _should_init_metrics_with_auth(): + return litellm.require_auth_for_metrics_endpoint + def prometheus_label_factory( supported_enum_labels: List[str], diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 2280e72e9b8..1c1b6f32c16 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -224,18 +224,9 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 litellm.callbacks = imported_list # type: ignore if "prometheus" in value: - if premium_user is not True: - verbose_proxy_logger.warning( - f"Prometheus metrics are only available for premium users. {CommonProxyErrors.not_premium_user.value}" - ) - from litellm.proxy.proxy_server import app + from litellm.integrations.prometheus import PrometheusLogger - verbose_proxy_logger.debug("Starting Prometheus Metrics on /metrics") - from prometheus_client import make_asgi_app - - # Add prometheus asgi middleware to route /metrics requests - metrics_app = make_asgi_app() - app.mount("/metrics", metrics_app) + PrometheusLogger._mount_metrics_endpoint(premium_user) else: litellm.callbacks = [ get_instance_fn( diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index fe8d73d26aa..788bf9642d5 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -4,12 +4,12 @@ model_list: model: openai/fake api_key: fake-key api_base: https://exampleopenaiendpoint-production.up.railway.app/ - -general_settings: - use_redis_transaction_buffer: true + - model_name: openai/gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: fake-key litellm_settings: - cache: True - cache_params: - type: redis - supported_call_types: [] \ No newline at end of file + require_auth_for_metrics_endpoint: true + callbacks: ["prometheus"] + service_callback: ["prometheus_system"] \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 100b0bf6dbd..1d2eab31c41 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1676,14 +1676,11 @@ class ProxyConfig: callback ) if "prometheus" in callback: - verbose_proxy_logger.debug( - "Starting Prometheus Metrics on /metrics" + from litellm.integrations.prometheus import ( + PrometheusLogger, ) - from prometheus_client import make_asgi_app - # Add prometheus asgi middleware to route /metrics requests - metrics_app = make_asgi_app() - app.mount("/metrics", metrics_app) + PrometheusLogger._mount_metrics_endpoint(premium_user) print( # noqa f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}" ) # noqa From f16c531002268a97d608e814e018f26cc4f8ba87 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 19:54:20 -0700 Subject: [PATCH 2/7] _mount_metrics_endpoint --- litellm/integrations/prometheus.py | 44 ------------------------------ 1 file changed, 44 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6cd413857bf..7c23ef86f4e 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1742,19 +1742,6 @@ class PrometheusLogger(CustomLogger): f"Prometheus metrics are only available for premium users. {CommonProxyErrors.not_premium_user.value}" ) - if PrometheusLogger._should_init_metrics_with_auth(): - PrometheusLogger._mount_metrics_endpoint_with_auth() - else: - # Mount metrics directly without authentication - PrometheusLogger._mount_metrics_endpoint_without_auth() - - @staticmethod - def _mount_metrics_endpoint_without_auth(): - from prometheus_client import make_asgi_app - - from litellm._logging import verbose_proxy_logger - from litellm.proxy.proxy_server import app - # Create metrics ASGI app metrics_app = make_asgi_app() @@ -1764,37 +1751,6 @@ class PrometheusLogger(CustomLogger): "Starting Prometheus Metrics on /metrics (no authentication)" ) - @staticmethod - def _mount_metrics_endpoint_with_auth(): - from fastapi import APIRouter, Depends - from prometheus_client import make_asgi_app - - from litellm._logging import verbose_proxy_logger - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.proxy_server import app - - # Create metrics ASGI app - metrics_app = make_asgi_app() - - # Create a router for authenticated metrics - metrics_router = APIRouter() - - # Add metrics endpoint with authentication - @metrics_router.get("/metrics", dependencies=[Depends(user_api_key_auth)]) - async def authenticated_metrics(): - verbose_proxy_logger.debug("Serving authenticated metrics endpoint") - return metrics_app - - # Mount the router to the app - app.include_router(metrics_router) - verbose_proxy_logger.debug( - "Starting Prometheus Metrics on /metrics with authentication" - ) - - @staticmethod - def _should_init_metrics_with_auth(): - return litellm.require_auth_for_metrics_endpoint - def prometheus_label_factory( supported_enum_labels: List[str], From c7523818b4abaf9dba269013d52a2e7637e64e09 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 20:27:17 -0700 Subject: [PATCH 3/7] PrometheusAuthMiddleware --- .../middleware/prometheus_auth_middleware.py | 39 +++++++++++++++++++ litellm/proxy/proxy_server.py | 2 + 2 files changed, 41 insertions(+) create mode 100644 litellm/proxy/middleware/prometheus_auth_middleware.py diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py new file mode 100644 index 00000000000..ae2481384b6 --- /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 1d2eab31c41..32fa4e8f52a 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 From 96ce5dbf7ddfe071d3dd3d4cb35efe0f1985db57 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 20:32:04 -0700 Subject: [PATCH 4/7] _should_run_auth_on_metrics_endpoint --- .../middleware/prometheus_auth_middleware.py | 28 +++++++++++++++---- litellm/proxy/proxy_config.yaml | 1 - 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index ae2481384b6..ca38a9d4452 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -5,6 +5,7 @@ from fastapi import Request from starlette.middleware.base import BaseHTTPMiddleware import litellm +from litellm.proxy._types import SpecialHeaders from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -13,12 +14,17 @@ class PrometheusAuthMiddleware(BaseHTTPMiddleware): # 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 + if self._should_run_auth_on_metrics_endpoint() is True: + try: + await user_api_key_auth( + request=request, + api_key=request.headers.get( + SpecialHeaders.openai_authorization.value + ) + or "", + ) + except Exception as e: + raise e # Process the request and get the response response = await call_next(request) @@ -36,4 +42,14 @@ class PrometheusAuthMiddleware(BaseHTTPMiddleware): @staticmethod def _should_run_auth_on_metrics_endpoint(): + """ + Returns True if auth should be run on the metrics endpoint + + False by default, set to True in proxy_config.yaml to enable + + ```yaml + litellm_settings: + require_auth_for_metrics_endpoint: true + ``` + """ return litellm.require_auth_for_metrics_endpoint diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 788bf9642d5..61950f55fe3 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -10,6 +10,5 @@ model_list: api_key: fake-key litellm_settings: - require_auth_for_metrics_endpoint: true callbacks: ["prometheus"] service_callback: ["prometheus_system"] \ No newline at end of file From 86b473d26708c2d3f3a13030efdf8fd637e97cfb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 20:37:17 -0700 Subject: [PATCH 5/7] allow adding auth on /metrics endpoint --- litellm/proxy/middleware/prometheus_auth_middleware.py | 6 +++++- litellm/proxy/proxy_config.yaml | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index ca38a9d4452..9e5a8ffcc8b 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -2,6 +2,7 @@ Prometheus Auth Middleware """ from fastapi import Request +from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware import litellm @@ -24,7 +25,10 @@ class PrometheusAuthMiddleware(BaseHTTPMiddleware): or "", ) except Exception as e: - raise e + return JSONResponse( + status_code=401, + content=f"Unauthorized access to metrics endpoint: {getattr(e, 'message', str(e))}", + ) # Process the request and get the response response = await call_next(request) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 61950f55fe3..788bf9642d5 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -10,5 +10,6 @@ model_list: api_key: fake-key litellm_settings: + require_auth_for_metrics_endpoint: true callbacks: ["prometheus"] service_callback: ["prometheus_system"] \ No newline at end of file From eaad3b24023760caf13ca267be1ea9bf21bb0c28 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 20:37:53 -0700 Subject: [PATCH 6/7] PrometheusAuthMiddleware --- .../proxy/middleware/prometheus_auth_middleware.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index 9e5a8ffcc8b..5d7913e5bf3 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -11,6 +11,19 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth class PrometheusAuthMiddleware(BaseHTTPMiddleware): + """ + Middleware to authenticate requests to the metrics endpoint + + By default, auth is not run on the metrics endpoint + + Enabled by setting the following in proxy_config.yaml: + + ```yaml + litellm_settings: + require_auth_for_metrics_endpoint: true + ``` + """ + async def dispatch(self, request: Request, call_next): # Check if this is a request to the metrics endpoint From fc4c453cb9447c628c8e3dd5eccfe83343e89726 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 4 Apr 2025 21:02:29 -0700 Subject: [PATCH 7/7] test_no_auth_metrics_when_disabled --- .../test_prometheus_auth_middleware.py | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py diff --git a/tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py b/tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py new file mode 100644 index 00000000000..b72ff75002b --- /dev/null +++ b/tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py @@ -0,0 +1,129 @@ +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +import pytest +from fastapi import FastAPI +from fastapi.responses import JSONResponse +from fastapi.testclient import TestClient + +import litellm +from litellm.proxy._types import SpecialHeaders +from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware + + +# Fake auth functions to simulate valid and invalid auth behavior. +async def fake_valid_auth(request, api_key): + # Simulate valid authentication: do nothing (i.e. pass) + return + + +async def fake_invalid_auth(request, api_key): + print("running fake invalid auth", request, api_key) + # Simulate invalid auth by raising an exception. + raise Exception("Invalid API key") + + +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + +@pytest.fixture +def app_with_middleware(): + """Create a FastAPI app with the PrometheusAuthMiddleware and dummy endpoints.""" + app = FastAPI() + # Add the PrometheusAuthMiddleware to the app. + app.add_middleware(PrometheusAuthMiddleware) + + @app.get("/metrics") + async def metrics(): + return {"msg": "metrics OK"} + + # Also allow /metrics/ (trailing slash) + @app.get("/metrics/") + async def metrics_slash(): + return {"msg": "metrics OK"} + + @app.get("/chat/completions") + async def chat(): + return {"msg": "chat completions OK"} + + @app.get("/embeddings") + async def embeddings(): + return {"msg": "embeddings OK"} + + return app + + +def test_valid_auth_metrics(app_with_middleware, monkeypatch): + """ + Test that a request to /metrics (and /metrics/) with valid auth headers passes. + """ + # Enable auth on metrics endpoints. + litellm.require_auth_for_metrics_endpoint = True + # Patch the auth function to simulate a valid authentication. + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + fake_valid_auth, + ) + + client = TestClient(app_with_middleware) + headers = {SpecialHeaders.openai_authorization.value: "valid"} + + # Test for /metrics (no trailing slash) + response = client.get("/metrics", headers=headers) + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"} + + # Test for /metrics/ (with trailing slash) + response = client.get("/metrics/", headers=headers) + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"} + + +def test_invalid_auth_metrics(app_with_middleware, monkeypatch): + """ + Test that a request to /metrics with invalid auth headers fails with a 401. + """ + litellm.require_auth_for_metrics_endpoint = True + # Patch the auth function to simulate a failed authentication. + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + fake_invalid_auth, + ) + + client = TestClient(app_with_middleware) + headers = {SpecialHeaders.openai_authorization.value: "invalid"} + + response = client.get("/metrics", headers=headers) + assert response.status_code == 401, response.text + assert "Unauthorized access to metrics endpoint" in response.text + + +def test_no_auth_metrics_when_disabled(app_with_middleware, monkeypatch): + """ + Test that when require_auth_for_metrics_endpoint is False, requests to /metrics + bypass the auth check. + """ + litellm.require_auth_for_metrics_endpoint = False + + # To ensure auth is not run, patch the auth function with one that will raise if called. + def should_not_be_called(*args, **kwargs): + raise Exception("Auth should not be called") + + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + should_not_be_called, + ) + + client = TestClient(app_with_middleware) + response = client.get("/metrics") + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"}