fix(prometheus.py): working e2e calls w/ userapikeymetadata

Author: krrishdholakia

enterprise/litellm_enterprise/integrations/prometheus.py

@@ -21,6 +21,7 @@
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth
 from litellm.types.integrations.prometheus import *
+from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name
 from litellm.types.utils import StandardLoggingPayload
 from litellm.utils import get_end_user_id_for_cost_tracking
 
@@ -2247,12 +2248,10 @@ def prometheus_label_factory(
 
     if enum_values.custom_metadata_labels is not None:
         for key, value in enum_values.custom_metadata_labels.items():
-            if key in supported_enum_labels:
-                filtered_labels[key] = value
-            else:
-                filtered_labels[key] = (
-                    "None"  # this happens for dynamically added metadata labels
-                )
+            # check sanitized key
+            sanitized_key = _sanitize_prometheus_label_name(key)
+            if sanitized_key in supported_enum_labels:
+                filtered_labels[sanitized_key] = value
 
     # Add custom tags if configured
     if enum_values.tags is not None:

litellm/proxy/_new_secret_config.yaml

@@ -6,4 +6,4 @@ model_list:
 
 litellm_settings:
   callbacks: ["prometheus"]
-  custom_prometheus_metadata_labels: ["metadata.initiative"]
+  custom_prometheus_metadata_labels: ["metadata.initiative", "metadata.business-unit"]

litellm/proxy/litellm_pre_call_utils.py

@@ -636,9 +636,11 @@ def add_management_endpoint_metadata_to_request_metadata(
                 + LiteLLM_ManagementEndpoint_MetadataFields
             ):
                 added_metadata[k] = v
-        data[_metadata_variable_name].setdefault(
-            "user_api_key_auth_metadata", {}
-        ).update(added_metadata)
+        if data[_metadata_variable_name].get("user_api_key_auth_metadata") is None:
+            data[_metadata_variable_name]["user_api_key_auth_metadata"] = {}
+        data[_metadata_variable_name]["user_api_key_auth_metadata"].update(
+            added_metadata
+        )
         return data
 
     @staticmethod

litellm/types/integrations/prometheus.py

@@ -426,13 +426,13 @@ class PrometheusMetricLabels:
 
     # Buffer monitoring metrics - these typically don't need additional labels
     litellm_pod_lock_manager_size: List[str] = []
-    
+
     litellm_in_memory_daily_spend_update_queue_size: List[str] = []
-    
+
     litellm_redis_daily_spend_update_queue_size: List[str] = []
-    
+
     litellm_in_memory_spend_update_queue_size: List[str] = []
-    
+
     litellm_redis_spend_update_queue_size: List[str] = []
 
     @staticmethod