From be70e659893e3c8afa9102ed634d40198fb305d1 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 27 Jun 2026 12:12:43 +0300 Subject: [PATCH] feat(prometheus): add per-team litellm_team_members_metric gauge Emit litellm_team_members_metric on every team member add and delete, labelled by team and team_alias and set to the team's authoritative member count. Because it is set from the current membership rather than incremented or decremented, it tracks the count up and down, never goes negative, and self-corrects on the next change after a proxy restart. Bulk member add is covered for free since it delegates to team_member_add, and the helper no-ops when the Prometheus callback is not registered. Resolves LIT-3082 --- litellm/integrations/prometheus.py | 23 ++++++ .../management_endpoints/team_endpoints.py | 23 ++++++ litellm/types/integrations/prometheus.py | 6 ++ .../test_prometheus_user_team_metrics.py | 77 +++++++++++++++++++ .../test_team_endpoints.py | 61 +++++++++++++++ 5 files changed, 190 insertions(+) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 18fb5d9491e1..f6ad46ebf39c 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -259,6 +259,13 @@ def __init__( ), ) + # Number of members in a team + self.litellm_team_members_metric = self._gauge_factory( + "litellm_team_members_metric", + "Number of members in a team", + labelnames=self.get_labels_for_metric("litellm_team_members_metric"), + ) + # Remaining Budget for Org self.litellm_remaining_org_budget_metric = self._gauge_factory( "litellm_remaining_org_budget_metric", @@ -3557,6 +3564,22 @@ def _set_team_budget_metrics( ) ) + def set_team_members_metric(self, team: LiteLLM_TeamTable) -> None: + """Set the team members gauge to the team's current member count.""" + enum_values = UserAPIKeyLabelValues( + team=team.team_id, + team_alias=team.team_alias or "", + ) + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_team_members_metric" + ), + enum_values=enum_values, + ) + self.litellm_team_members_metric.labels(**_labels).set( + len(team.members_with_roles) + ) + async def _set_org_budget_metrics_after_api_request( self, org_id: Optional[str], diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f533774ce180..ac7d92aaa422 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -23,6 +23,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( BlockTeamRequest, @@ -2445,6 +2446,23 @@ async def _add_team_members_to_team( return updated_team, updated_users, updated_team_memberships +def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: + """Update the Prometheus team members gauge after a membership change. + + No-ops when the Prometheus callback is not registered, and never lets a + metric failure break the team add/delete request. + """ + prometheus_logger = PrometheusLogger.get_instance() + if prometheus_logger is None: + return + try: + prometheus_logger.set_team_members_metric(team) + except Exception as e: + verbose_proxy_logger.debug( + "Prometheus: failed to emit team members metric: %s", str(e) + ) + + async def _validate_and_populate_member_user_info( member: Member, prisma_client: PrismaClient, @@ -2665,6 +2683,9 @@ async def team_member_add( raise HTTPException( status_code=404, detail={"error": f"Team with id {data.team_id} not found"} ) + + _emit_team_members_metric(complete_team_data) + return TeamAddMemberResponse( **updated_team.model_dump(), updated_users=updated_users, @@ -2792,6 +2813,8 @@ async def team_member_delete( data={"members_with_roles": json.dumps(_db_new_team_members)}, # type: ignore ) + _emit_team_members_metric(existing_team_row) + ## DELETE TEAM ID from USER ROW, IF EXISTS ## # get user row key_val = {} diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index ba427de41a2d..52afc889ae6d 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -215,6 +215,7 @@ class UserAPIKeyLabelNames(Enum): "litellm_remaining_team_budget_metric", "litellm_team_max_budget_metric", "litellm_team_budget_remaining_hours_metric", + "litellm_team_members_metric", "litellm_remaining_org_budget_metric", "litellm_org_max_budget_metric", "litellm_org_budget_remaining_hours_metric", @@ -533,6 +534,11 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, ] + litellm_team_members_metric = [ + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + ] + litellm_remaining_org_budget_metric = [ UserAPIKeyLabelNames.ORG_ID.value, UserAPIKeyLabelNames.ORG_ALIAS.value, 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 361ab7332f80..90a9d1fcceba 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -927,3 +927,80 @@ def test_custom_latency_buckets(): REGISTRY.unregister(collector) except Exception: pass + + +class TestSetTeamMembersMetric: + """litellm_team_members_metric tracks the current member count per team.""" + + def _gauge_value(self, team_id, team_alias): + return REGISTRY.get_sample_value( + "litellm_team_members_metric", + {"team": team_id, "team_alias": team_alias}, + ) + + def test_metric_initialized(self, prometheus_logger): + assert hasattr(prometheus_logger, "litellm_team_members_metric") + assert prometheus_logger.litellm_team_members_metric is not None + + @pytest.mark.parametrize("count", [0, 1, 3, 7]) + def test_sets_gauge_to_member_count(self, prometheus_logger, count): + from litellm.proxy._types import LiteLLM_TeamTable, Member + + team = LiteLLM_TeamTable( + team_id="team-a", + team_alias="Acme", + members_with_roles=[ + Member(user_id=f"u{i}", role="user") for i in range(count) + ], + ) + prometheus_logger.set_team_members_metric(team) + assert self._gauge_value("team-a", "Acme") == float(count) + + def test_gauge_reflects_latest_count_not_delta(self, prometheus_logger): + """Re-emitting overwrites with the authoritative count (set, not inc/dec).""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + + members = [Member(user_id=f"u{i}", role="user") for i in range(4)] + team = LiteLLM_TeamTable( + team_id="team-b", team_alias="Beta", members_with_roles=members + ) + prometheus_logger.set_team_members_metric(team) + assert self._gauge_value("team-b", "Beta") == 4.0 + + # Drop two members and re-emit: gauge must read 2, not 4 and not -2. + team.members_with_roles = members[:2] + prometheus_logger.set_team_members_metric(team) + assert self._gauge_value("team-b", "Beta") == 2.0 + + def test_none_alias_falls_back_to_empty_string(self, prometheus_logger): + from litellm.proxy._types import LiteLLM_TeamTable, Member + + team = LiteLLM_TeamTable( + team_id="team-c", + team_alias=None, + members_with_roles=[Member(user_id="solo", role="admin")], + ) + prometheus_logger.set_team_members_metric(team) + assert self._gauge_value("team-c", "") == 1.0 + + def test_teams_isolated_by_label(self, prometheus_logger): + from litellm.proxy._types import LiteLLM_TeamTable, Member + + team_one = LiteLLM_TeamTable( + team_id="team-1", + team_alias="One", + members_with_roles=[Member(user_id="a", role="user")], + ) + team_two = LiteLLM_TeamTable( + team_id="team-2", + team_alias="Two", + members_with_roles=[ + Member(user_id="b", role="user"), + Member(user_id="c", role="user"), + Member(user_id="d", role="user"), + ], + ) + prometheus_logger.set_team_members_metric(team_one) + prometheus_logger.set_team_members_metric(team_two) + assert self._gauge_value("team-1", "One") == 1.0 + assert self._gauge_value("team-2", "Two") == 3.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ecce92e7324a..5f3974b46fbf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9434,3 +9434,64 @@ async def test_team_info_forwards_key_limit_to_get_data(): ) assert mock_prisma.get_data.await_args.kwargs["limit"] == 7 + + +class TestEmitTeamMembersMetric: + """The _emit_team_members_metric seam between the team handlers and Prometheus.""" + + @pytest.fixture + def restore_callbacks(self): + import litellm + + original = litellm.callbacks + yield + litellm.callbacks = original + + def _team(self, member_count): + return LiteLLM_TeamTable( + team_id="team-x", + team_alias="X", + members_with_roles=[ + Member(user_id=f"u{i}", role="user") for i in range(member_count) + ], + ) + + def test_emits_with_team_when_logger_registered(self, restore_callbacks): + import litellm + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric, + ) + + fake_logger = MagicMock(spec=PrometheusLogger) + litellm.callbacks = [fake_logger] + + team = self._team(3) + _emit_team_members_metric(team) + + fake_logger.set_team_members_metric.assert_called_once_with(team) + + def test_noop_when_no_logger_registered(self, restore_callbacks): + import litellm + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric, + ) + + litellm.callbacks = [] + # Must not raise when Prometheus is not enabled. + _emit_team_members_metric(self._team(2)) + + def test_metric_failure_does_not_break_request(self, restore_callbacks): + import litellm + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric, + ) + + fake_logger = MagicMock(spec=PrometheusLogger) + fake_logger.set_team_members_metric.side_effect = Exception("boom") + litellm.callbacks = [fake_logger] + + # A metric failure must be swallowed, not propagated to the handler. + _emit_team_members_metric(self._team(1)) + fake_logger.set_team_members_metric.assert_called_once()