diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 24597c02ea2..41a4d026fe1 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -69,6 +69,15 @@ else: _DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT = 5.0 +# Tiers a caller may name in a request, across the providers that accept the +# parameter: OpenAI ("auto", "default", "flex", "priority", "scale"), Bedrock and +# Groq (subsets of those), Anthropic ("auto", "standard_only") and Vertex, which +# maps "default" to "standard". Used to bound the caller-controlled fallback in +# ``get_service_tier_from_standard_logging_payload``. +KNOWN_REQUEST_SERVICE_TIERS = frozenset( + {"auto", "batch", "default", "flex", "priority", "scale", "standard", "standard_only"} +) + def _get_budget_metrics_per_request_timeout() -> float: raw = os.getenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT") @@ -1245,6 +1254,7 @@ class PrometheusLogger(CustomLogger): client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), user_agent=standard_logging_payload["metadata"].get("user_agent"), stream=(str(standard_logging_payload.get("stream")) if litellm.prometheus_emit_stream_label else None), + service_tier=get_service_tier_from_standard_logging_payload(standard_logging_payload), ) if user_api_key is not None and isinstance(user_api_key, str) and user_api_key.startswith("sk-"): @@ -4098,6 +4108,44 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]: return result +def get_service_tier_from_standard_logging_payload( + standard_logging_payload: StandardLoggingPayload, +) -> str | None: + """ + Resolve the service tier a request ran on, for the ``service_tier`` label. + + The tier the provider actually served wins over the tier the caller asked for, + so latency and spend stay segmentable when the request said ``auto`` and the + provider picked the concrete tier. Providers report the served tier either at + the top level of the response (OpenAI, Bedrock, Groq) or on the usage object + (Anthropic). + + Streaming responses carry no served tier, so the requested tier is the + fallback. That value is caller-controlled and survives param mapping even + where the provider then ignores it (Bedrock and Groq accept the request and + drop an unrecognized tier), so it is only labelled when it names a known + tier; otherwise one caller could mint a Prometheus series per string. Values + the provider itself reports are not caller-controlled and stay unrestricted, + so a tier a provider adds later is still labelled correctly. + """ + response = standard_logging_payload.get("response") + usage_object = standard_logging_payload.get("metadata", {}).get("usage_object") + + served_candidates: tuple[object, ...] = ( + response.get("service_tier") if isinstance(response, dict) else None, + usage_object.get("service_tier") if isinstance(usage_object, dict) else None, + ) + served_tier = next((tier for tier in served_candidates if isinstance(tier, str) and tier), None) + if served_tier is not None: + return served_tier + + model_parameters = standard_logging_payload.get("model_parameters") + requested_tier = model_parameters.get("service_tier") if isinstance(model_parameters, dict) else None + if isinstance(requested_tier, str) and requested_tier in KNOWN_REQUEST_SERVICE_TIERS: + return requested_tier + return None + + def _get_combined_custom_metadata_from_standard_logging_payload( standard_logging_payload: Optional[dict], ) -> Dict[str, Any]: diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 318ba1f5956..905efc84b5a 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -190,6 +190,7 @@ class UserAPIKeyLabelNames(Enum): ORG_ALIAS = "org_alias" MCP_TOOL_NAME = "mcp_tool_name" MCP_SERVER_NAME = "mcp_server_name" + SERVICE_TIER = "service_tier" DEFINED_PROMETHEUS_METRICS = Literal[ @@ -286,6 +287,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.MODEL_ID.value, UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.SERVICE_TIER.value, ] litellm_llm_api_time_to_first_token_metric = [ @@ -299,6 +301,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER.value, UserAPIKeyLabelNames.MODEL_ID.value, UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.SERVICE_TIER.value, ] litellm_request_total_latency_metric = [ @@ -312,6 +315,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.SERVICE_TIER.value, ] litellm_request_queue_time_seconds = [ @@ -453,6 +457,7 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.REQUESTED_MODEL.value, UserAPIKeyLabelNames.MODEL_ID.value, UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.SERVICE_TIER.value, ] litellm_input_tokens_metric = [ @@ -878,6 +883,7 @@ class UserAPIKeyLabelValues: org_alias: Optional[str] = None mcp_tool_name: Optional[str] = None mcp_server_name: Optional[str] = None + service_tier: Optional[str] = None # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 1b1ce3e0f2d..9acb87750e9 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -428,6 +428,7 @@ def test_set_latency_metrics(prometheus_logger): model="gpt-5-mini", model_id="model-123", api_provider="openai", + service_tier=None, ) prometheus_logger.litellm_llm_api_time_to_first_token_metric.labels().observe.assert_called_once_with( 0.5 @@ -447,6 +448,7 @@ def test_set_latency_metrics(prometheus_logger): model="gpt-5-mini", model_id="model-123", api_provider="openai", + service_tier=None, ) prometheus_logger.litellm_llm_api_latency_metric.labels().observe.assert_called_once_with( 1.5 @@ -466,6 +468,7 @@ def test_set_latency_metrics(prometheus_logger): model="gpt-5-mini", model_id="model-123", api_provider="openai", + service_tier=None, ) prometheus_logger.litellm_request_total_latency_metric.labels().observe.assert_called_once_with( 2.0 @@ -634,6 +637,7 @@ def test_increment_top_level_request_and_spend_metrics(prometheus_logger): client_ip=None, user_agent=None, requested_model=None, + service_tier=None, ) prometheus_logger.litellm_spend_metric.labels().inc.assert_called_once_with(0.1) diff --git a/tests/test_litellm/integrations/test_prometheus_service_tier_label.py b/tests/test_litellm/integrations/test_prometheus_service_tier_label.py new file mode 100644 index 00000000000..9d702b19c6c --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_service_tier_label.py @@ -0,0 +1,232 @@ +""" +Unit tests for the service_tier Prometheus label on latency and spend metrics. + +Covers the label being declared on the metrics that carry it, the precedence +between the tier a provider served and the tier a caller requested, and the +end-to-end emit wiring through async_log_success_event. + +Run with: + uv run pytest tests/test_litellm/integrations/test_prometheus_service_tier_label.py -v +""" + +import datetime + +import pytest + +from litellm.integrations.prometheus import ( + KNOWN_REQUEST_SERVICE_TIERS, + PrometheusLogger, + get_service_tier_from_standard_logging_payload, +) +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, +) + +SERVICE_TIER_METRICS = [ + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", + "litellm_spend_metric", +] + + +def _clear_prometheus_registry() -> None: + from prometheus_client import REGISTRY + + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _collected_samples(metric_name: str): + from prometheus_client import REGISTRY + + return [sample for metric in REGISTRY.collect() for sample in metric.samples if sample.name == metric_name] + + +def _standard_logging_payload( + response: object = None, + usage_object: object = None, + model_parameters: object = None, +) -> dict: + return { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": 30, + "prompt_tokens": 20, + "completion_tokens": 10, + "startTime": 1.0, + "endTime": 2.0, + "completionStartTime": 1.5, + "model": "gpt-4o-mini", + "model_id": "model-123", + "model_group": "gpt-4o-mini", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "request_tags": [], + "end_user": None, + "cache_hit": False, + "stream": True, + "response": response, + "model_parameters": model_parameters, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": usage_object, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + } + + +def test_service_tier_label_declared_on_latency_and_spend_metrics(): + for metric_name in SERVICE_TIER_METRICS: + labels = PrometheusMetricLabels.get_labels(metric_name) + assert UserAPIKeyLabelNames.SERVICE_TIER.value in labels, f"{metric_name} should carry the service_tier label" + + +def test_user_api_key_label_values_carries_service_tier(): + values = UserAPIKeyLabelValues(service_tier="flex") + + assert values.service_tier == "flex" + assert values.model_dump()["service_tier"] == "flex" + assert UserAPIKeyLabelValues().service_tier is None + + +def test_served_tier_wins_over_requested_tier(): + payload = _standard_logging_payload( + response={"service_tier": "default"}, + model_parameters={"service_tier": "auto"}, + ) + + assert get_service_tier_from_standard_logging_payload(payload) == "default" + + +def test_usage_object_tier_used_when_response_has_none(): + payload = _standard_logging_payload( + response={"id": "chatcmpl-1"}, + usage_object={"prompt_tokens": 1, "service_tier": "standard"}, + model_parameters={"service_tier": "auto"}, + ) + + assert get_service_tier_from_standard_logging_payload(payload) == "standard" + + +def test_requested_tier_used_when_no_served_tier(): + payload = _standard_logging_payload( + response={"id": "chatcmpl-1"}, + usage_object={"prompt_tokens": 1}, + model_parameters={"service_tier": "flex"}, + ) + + assert get_service_tier_from_standard_logging_payload(payload) == "flex" + + +def test_unrecognized_requested_tier_is_not_labelled(): + """ + A caller-supplied tier survives param mapping even where the provider then + ignores it (Bedrock and Groq drop an unrecognized tier and still answer), so + labelling it verbatim would let one caller mint a series per string. + """ + payload = _standard_logging_payload(model_parameters={"service_tier": "attacker-controlled-a1b2c3"}) + + assert get_service_tier_from_standard_logging_payload(payload) is None + + +def test_unrecognized_served_tier_is_labelled(): + """ + The tier a provider reports is not caller-controlled, so a tier added by a + provider after this release still gets labelled instead of being dropped. + """ + payload = _standard_logging_payload( + response={"service_tier": "tier-added-by-provider-later"}, + model_parameters={"service_tier": "auto"}, + ) + + assert get_service_tier_from_standard_logging_payload(payload) == "tier-added-by-provider-later" + + +@pytest.mark.parametrize("tier", sorted(KNOWN_REQUEST_SERVICE_TIERS)) +def test_every_known_requested_tier_is_labelled(tier): + payload = _standard_logging_payload(model_parameters={"service_tier": tier}) + + assert get_service_tier_from_standard_logging_payload(payload) == tier + + +@pytest.mark.parametrize( + "response, usage_object, model_parameters", + [ + (None, None, None), + ({"service_tier": None}, {"service_tier": ""}, {"service_tier": None}), + ("redacted-by-litellm", None, {}), + ({"service_tier": 1}, None, None), + ], +) +def test_no_tier_resolves_to_none(response, usage_object, model_parameters): + payload = _standard_logging_payload( + response=response, + usage_object=usage_object, + model_parameters=model_parameters, + ) + + assert get_service_tier_from_standard_logging_payload(payload) is None + + +@pytest.mark.asyncio +async def test_success_event_emits_service_tier_on_latency_and_spend_metrics(): + """ + End-to-end emit wiring. + + Drives the real logger with a payload whose response was served on the flex + tier and asserts every latency histogram and the spend counter carries + service_tier="flex". Fails if the label is dropped from a metric's label list + or if the value is not populated on the success path. + """ + payload = _standard_logging_payload( + response={"service_tier": "flex"}, + model_parameters={"service_tier": "auto"}, + ) + now = datetime.datetime.now() + kwargs = { + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {}}, + "standard_logging_object": payload, + "stream": True, + "start_time": now - datetime.timedelta(seconds=3), + "api_call_start_time": now - datetime.timedelta(seconds=2), + "completion_start_time": now - datetime.timedelta(seconds=1), + "end_time": now, + } + + _clear_prometheus_registry() + try: + logger = PrometheusLogger() + await logger.async_log_success_event(kwargs, None, now, now) + + for metric_name in ( + "litellm_request_total_latency_metric_bucket", + "litellm_llm_api_latency_metric_bucket", + "litellm_llm_api_time_to_first_token_metric_bucket", + "litellm_spend_metric_total", + ): + samples = _collected_samples(metric_name) + assert samples, f"expected {metric_name} to be emitted" + assert all(sample.labels.get("service_tier") == "flex" for sample in samples), ( + f"{metric_name} must carry service_tier=flex, got " + f"{sorted({sample.labels.get('service_tier') for sample in samples})}" + ) + finally: + _clear_prometheus_registry()