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..7c23ef86f4e 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1721,6 +1721,36 @@ 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}" + ) + + # 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)" + ) + 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/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py new file mode 100644 index 00000000000..5d7913e5bf3 --- /dev/null +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -0,0 +1,72 @@ +""" +Prometheus Auth Middleware +""" +from fastapi import Request +from fastapi.responses import JSONResponse +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 + + +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 + + if self._is_prometheus_metrics_endpoint(request): + 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: + 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) + + 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(): + """ + 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 56eb2ca39f2..da8bc383b1e 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -4,7 +4,13 @@ model_list: model: openai/fake api_key: fake-key api_base: https://exampleopenaiendpoint-production.up.railway.app/ + - model_name: openai/gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: fake-key litellm_settings: + 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 acc6b6175e9..e1982a3ca0b 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 @@ -1676,14 +1678,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 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"}