fix(prometheus): retain deployment attribution on proxy failures (#49717)

(cherry picked from commit 1563776cd8)

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

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