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()