From b2e708d5aecdcdb5a19c84aa5f211844e427522b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 27 Jun 2026 22:19:50 +0300 Subject: [PATCH] feat(prometheus): add per-team litellm_team_members_metric gauge (#31506) 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 18fb5d9491e..f6ad46ebf39 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -259,6 +259,13 @@ class PrometheusLogger(CustomLogger): ), ) + # 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 @@ class PrometheusLogger(CustomLogger): ) ) + 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 f533774ce18..ac7d92aaa42 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -23,6 +23,7 @@ from pydantic import BaseModel 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 ba427de41a2..52afc889ae6 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -215,6 +215,7 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "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 361ab7332f8..90a9d1fcceb 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 ecce92e7324..5f3974b46fb 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()