diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fb5fc253ae4..c395987695b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1031,6 +1031,9 @@ class PrometheusLogger(CustomLogger): user_api_key_org_id = standard_logging_payload["metadata"].get( "user_api_key_org_id" ) + user_api_key_org_alias = standard_logging_payload["metadata"].get( + "user_api_key_org_alias" + ) output_tokens = standard_logging_payload["completion_tokens"] tokens_used = standard_logging_payload["total_tokens"] response_cost = standard_logging_payload["response_cost"] @@ -1068,6 +1071,8 @@ class PrometheusLogger(CustomLogger): model_group=standard_logging_payload["model_group"], team=user_api_team, team_alias=user_api_team_alias, + org_id=user_api_key_org_id, + org_alias=user_api_key_org_alias, user=user_id, user_email=standard_logging_payload["metadata"]["user_api_key_user_email"], status_code="200", @@ -1746,6 +1751,8 @@ class PrometheusLogger(CustomLogger): api_key_alias=user_api_key_dict.key_alias, team=user_api_key_dict.team_id, team_alias=user_api_key_dict.team_alias, + org_id=user_api_key_dict.org_id, + org_alias=user_api_key_dict.organization_alias, requested_model=request_data.get("model", ""), status_code=str(status_code), exception_status=str(status_code), diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 441b3b836a1..70f93619314 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2419,6 +2419,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): end_user_model_max_budget: Optional[dict] = None # Organization Params + organization_alias: Optional[str] = None organization_max_budget: Optional[float] = None organization_tpm_limit: Optional[int] = None organization_rpm_limit: Optional[int] = None diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3ed96c163af..ece72b1060e 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -684,6 +684,7 @@ class LiteLLMProxyRequestSetup: user_api_key_project_alias=user_api_key_dict.project_alias, user_api_key_user_id=user_api_key_dict.user_id, user_api_key_org_id=user_api_key_dict.org_id, + user_api_key_org_alias=user_api_key_dict.organization_alias, user_api_key_team_alias=user_api_key_dict.team_alias, user_api_key_end_user_id=user_api_key_dict.end_user_id, user_api_key_user_email=user_api_key_dict.user_email, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9738ae4f1a2..4cc427cc9a4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7128,6 +7128,7 @@ async def chat_completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias if ( hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None @@ -7302,6 +7303,7 @@ async def completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias if ( hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None @@ -7544,6 +7546,7 @@ async def embeddings( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias if ( hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec98cfd4d1e..db5fcd98e5e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3004,6 +3004,7 @@ class PrismaClient: b.model_max_budget as litellm_budget_table_model_max_budget, b.soft_budget as litellm_budget_table_soft_budget, o.metadata as organization_metadata, + o.organization_alias as organization_alias, b2.max_budget as organization_max_budget, b2.tpm_limit as organization_tpm_limit, b2.rpm_limit as organization_rpm_limit diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 0d1501664b9..5f1aa9fb2ce 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,7 +1,7 @@ import re from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, List, Literal, Optional, Tuple +from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple from pydantic import BaseModel, Field, field_validator from typing_extensions import Annotated @@ -665,6 +665,24 @@ class PrometheusMetricLabels: litellm_cache_misses_metric = _cache_metric_labels litellm_cached_tokens_metric = _cache_metric_labels + # Metrics whose emission paths supply org context (used by get_labels) + _org_label_metrics: ClassVar[frozenset] = frozenset( + { + "litellm_llm_api_latency_metric", + "litellm_llm_api_time_to_first_token_metric", + "litellm_request_total_latency_metric", + "litellm_request_queue_time_seconds", + "litellm_proxy_total_requests_metric", + "litellm_proxy_failed_requests_metric", + "litellm_deployment_latency_per_output_token", + "litellm_requests_metric", + "litellm_spend_metric", + "litellm_input_tokens_metric", + "litellm_total_tokens_metric", + "litellm_output_tokens_metric", + } + ) + # Managed batch metrics _batch_user_labels = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -731,6 +749,14 @@ class PrometheusMetricLabels: ): custom_labels.append(UserAPIKeyLabelNames.STREAM.value) + if label_name in PrometheusMetricLabels._org_label_metrics: + for label in [ + UserAPIKeyLabelNames.ORG_ID.value, + UserAPIKeyLabelNames.ORG_ALIAS.value, + ]: + if label not in default_labels and label not in custom_labels: + custom_labels.append(label) + return default_labels + custom_labels diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3f6e6e5aa5a..cd5806b3ab7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2507,6 +2507,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_max_budget: Optional[float] user_api_key_budget_reset_at: Optional[str] user_api_key_org_id: Optional[str] + user_api_key_org_alias: Optional[str] user_api_key_team_id: Optional[str] user_api_key_project_id: Optional[str] user_api_key_project_alias: Optional[str] diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py index 4bfa3a581e3..48d9cbd1bb1 100644 --- a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -118,6 +118,8 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent(): "user_api_key_alias": "alias_1", "user_api_key_team_id": "team_1", "user_api_key_team_alias": "team_alias_1", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, "user_api_key_user_email": "test@example.com", "user_api_key_request_route": "/chat/completions", "requester_ip_address": "192.168.1.1", diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 9bcf08fdd71..6b65f444046 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -525,6 +525,61 @@ async def test_set_user_budget_metrics_after_api_request_inf_when_genuinely_no_b ) +def test_per_request_metrics_emit_all_identity_labels(prometheus_logger): + """Verify org labels appear when flag is on and are absent when flag is off.""" + import litellm + from litellm.types.integrations.prometheus import UserAPIKeyLabelValues + + prometheus_logger.litellm_requests_metric = MagicMock() + prometheus_logger.litellm_spend_metric = MagicMock() + + enum_values = UserAPIKeyLabelValues( + hashed_api_key="hashed-key", + api_key_alias="my-key", + model="gpt-4", + team="team-abc", + team_alias="my-team", + org_id="org-abc", + org_alias="my-org", + user="user-1", + ) + + common_kwargs = dict( + end_user_id=None, + user_api_key="hashed-key", + user_api_key_alias="my-key", + model="gpt-4", + user_api_team="team-abc", + user_api_team_alias="my-team", + user_id="user-1", + response_cost=0.001, + enum_values=enum_values, + ) + + try: + # org labels are always included in per-request metrics + prometheus_logger._increment_top_level_request_and_spend_metrics(**common_kwargs) + label_kwargs = prometheus_logger.litellm_requests_metric.labels.call_args.kwargs + assert label_kwargs["org_id"] == "org-abc" + assert label_kwargs["org_alias"] == "my-org" + assert label_kwargs["team"] == "team-abc" + assert label_kwargs["user"] == "user-1" + + # Metrics not in the org-emission list must NOT get org labels + from litellm.types.integrations.prometheus import PrometheusMetricLabels + for metric in ("litellm_remaining_api_key_budget_metric", "litellm_remaining_team_budget_metric"): + labels = PrometheusMetricLabels.get_labels(metric) + assert "org_id" not in labels, f"{metric} should not have org_id" + assert "org_alias" not in labels, f"{metric} should not have org_alias" + + # org_id in custom_prometheus_metadata_labels must not produce duplicate labels + litellm.custom_prometheus_metadata_labels = ["org_id"] + labels = PrometheusMetricLabels.get_labels("litellm_requests_metric") + assert labels.count("org_id") == 1 + finally: + litellm.custom_prometheus_metadata_labels = [] + + # --------------------------------------------------------------------------- # Org budget metric tests # ---------------------------------------------------------------------------