diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 97a1073d7c0..21b17af88a5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -259,6 +259,15 @@ _POSITIVE_SERIES_CAP: Final[TypeAdapter[int]] = TypeAdapter(Annotated[int, Field _POSITIVE_SERIES_TTL: Final[TypeAdapter[float]] = TypeAdapter(Annotated[float, Field(gt=0)]) _SERIES_CLEANUP_INTERVAL: Final[TypeAdapter[float]] = TypeAdapter(Annotated[float, Field(ge=0)]) _DEFAULT_SERIES_CLEANUP_INTERVAL_SECONDS: Final = 60.0 +_LOGGED_FIELDS: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_NO_LOGGED_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _logged_fields(value: object) -> Mapping[str, object]: + try: + return _LOGGED_FIELDS.validate_python(value) + except ValidationError: + return _NO_LOGGED_FIELDS def _number_or_none(value: object, limit: TypeAdapter[_SeriesLimitT]) -> _SeriesLimitT | None: @@ -2847,7 +2856,7 @@ class PrometheusLogger(CustomLogger): async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, @@ -2870,10 +2879,14 @@ class PrometheusLogger(CustomLogger): try: _tags: Final = StandardLoggingPayloadSetup.get_request_tags( litellm_params=request_data, - proxy_server_request=request_data.get("proxy_server_request", {}), + proxy_server_request=_logged_fields(request_data.get("proxy_server_request")), + ) + _metadata: Final = _logged_fields(request_data.get("metadata")) + model_id: Final = ( + _logged_fields(_metadata.get("model_info")).get("id") + or _logged_fields(request_data.get("model_info")).get("id") + or _logged_fields(request_data.get("standard_logging_object")).get("model_id") ) - _metadata: Final = request_data.get("metadata", {}) or {} - model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) api_provider: Final = self._extract_api_provider_from_request_data( request_data diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index ead0ce8ecff..5d6a11c489f 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -2,7 +2,10 @@ Unit tests for prometheus metric labels configuration """ +from typing import Final + import pytest +from pydantic import JsonValue from litellm.types.integrations.prometheus import ( PrometheusMetricLabels, @@ -716,6 +719,47 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() _clear_prometheus_registry() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_id", "request_id", "logged_id", "expected"), + ( + (None, None, "logged-deployment", "logged-deployment"), + ("", None, "logged-deployment", "logged-deployment"), + ("metadata-deployment", "request-deployment", "logged-deployment", "metadata-deployment"), + (None, "request-deployment", "logged-deployment", "request-deployment"), + (None, None, None, "None"), + ), +) +async def test_failure_counters_keep_logged_deployment_attribution( + metadata_id: str | None, request_id: str | None, logged_id: str | None, expected: str +) -> None: + from fastapi import HTTPException + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + request_data: Final[dict[str, JsonValue]] = { + "model": "metric-test-alias", + "metadata": {"model_info": {"id": metadata_id}}, + "model_info": {"id": request_id}, + "standard_logging_object": {"model_id": logged_id, "custom_llm_provider": "openai"}, + } + _clear_prometheus_registry() + try: + await PrometheusLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="guardrail rejected"), + user_api_key_dict=UserAPIKeyAuth(token="metric-test-token"), + ) + for metric_name in ("litellm_proxy_failed_requests_metric_total", "litellm_proxy_total_requests_metric_total"): + samples: Final = _collected_samples(metric_name) + assert [(sample.labels["model_id"], sample.value) for sample in samples] == [ + (expected, 1.0) + ] + finally: + _clear_prometheus_registry() + + @pytest.mark.asyncio async def test_failure_hook_labels_a_cli_session_with_the_per_user_alias_not_the_login_token(): from litellm.integrations.prometheus import PrometheusLogger