scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 213

Feat - Add organization into the metrics metadata for org_id & org_alias (#24440)

Browse source 
* Add org_id and org_alias label names to Prometheus metric definitions

* Add user_api_key_org_alias to StandardLoggingUserAPIKeyMetadata

* Populate user_api_key_org_alias in pre-call metadata

* Pass org_id and org_alias into per-request Prometheus metric labels

* Add test for org labels on per-request Prometheus metrics

* chore: resolve test mockdata

* Address review: populate org_alias from DB view, add feature flag, use .get() for org metadata

* Add org labels to failure path and verify flag behavior in test

* Fix test: build flag-off enum_values without org fields

* Gate org labels behind feature flag in get_labels() instead of static metric lists

* Scope org label injection to metrics that carry team context, remove orphaned budget label defs, add test teardown

* Use explicit metric allowlist for org label injection instead of team heuristic

* Fix duplicate org label guard, move _org_label_metrics to class constant

* Reset custom_prometheus_metadata_labels after duplicate label assertion

* fix: emit org labels by default, remove flag, fix missing org_alias in all metadata paths

* fix: emit org labels by default, no opt-in flag required

* fix: write org_alias to metadata unconditionally in proxy_server.py
This commit is contained in:
Josh 2026-04-03 00:46:37 -04:00 • committed by Sameer Kankute
parent 0873145727
commit 0f48f793ff
No known key found for this signature in database
9 changed files with 98 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/integrations/prometheus.py
View file

@ -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),

1
litellm/proxy/_types.py
View file

@ -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

1
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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,

3
litellm/proxy/proxy_server.py
View file

@ -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

1
litellm/proxy/utils.py
View file

@ -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

28
litellm/types/integrations/prometheus.py
View file

@ -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

1
litellm/types/utils.py
View file

@ -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]

2
tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py
View file

@ -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",

55
tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
View file

@ -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
# ---------------------------------------------------------------------------