diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 30af0dcb8ed..68340f55d99 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -477,6 +477,13 @@ class PrometheusLogger(CustomLogger): labelnames=[], ) + # Per-team member count (incremented on add, decremented on remove) + self.litellm_team_members_metric = self._gauge_factory( + "litellm_team_members_metric", + "Current number of members per team. Incremented when a member is added, decremented when a member is removed.", + labelnames=self.get_labels_for_metric("litellm_team_members_metric"), + ) + ######################################## # Managed Batch Metrics ######################################## @@ -3268,6 +3275,57 @@ class PrometheusLogger(CustomLogger): ) ) + def _team_members_metric_labels( + self, + team_id: Optional[str], + team_alias: Optional[str], + ) -> dict: + """Build the label dict for litellm_team_members_metric.""" + enum_values = UserAPIKeyLabelValues( + team=team_id or "", + team_alias=team_alias or "", + ) + return prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_team_members_metric" + ), + enum_values=enum_values, + ) + + def increment_team_members_metric( + self, + team_id: Optional[str], + team_alias: Optional[str], + amount: float = 1.0, + ) -> None: + """Increment the team-members gauge by ``amount`` (default +1).""" + try: + labels = self._team_members_metric_labels( + team_id=team_id, team_alias=team_alias + ) + self.litellm_team_members_metric.labels(**labels).inc(amount) + except Exception as e: + verbose_logger.exception( + "Error incrementing litellm_team_members_metric: %s", str(e) + ) + + def decrement_team_members_metric( + self, + team_id: Optional[str], + team_alias: Optional[str], + amount: float = 1.0, + ) -> None: + """Decrement the team-members gauge by ``amount`` (default -1).""" + try: + labels = self._team_members_metric_labels( + team_id=team_id, team_alias=team_alias + ) + self.litellm_team_members_metric.labels(**labels).dec(amount) + except Exception as e: + verbose_logger.exception( + "Error decrementing litellm_team_members_metric: %s", str(e) + ) + 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 35e3d196e9e..c954568bf1e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -100,6 +100,7 @@ from litellm.proxy.management_helpers.utils import ( ) from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy from litellm.router import Router +from litellm.integrations.prometheus import PrometheusLogger from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -119,6 +120,32 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( router = APIRouter() +def _emit_team_members_metric_delta( + team_id: Optional[str], + team_alias: Optional[str], + delta: int, +) -> None: + """Emit a delta to the per-team Prometheus team-members gauge. + + Positive ``delta`` increments the gauge (member added), negative decrements + (member removed). Silently no-ops when the Prometheus callback is not + registered or when ``delta`` is zero. + """ + if delta == 0: + return + prometheus_logger = PrometheusLogger.get_instance() + if prometheus_logger is None: + return + if delta > 0: + prometheus_logger.increment_team_members_metric( + team_id=team_id, team_alias=team_alias, amount=float(delta) + ) + else: + prometheus_logger.decrement_team_members_metric( + team_id=team_id, team_alias=team_alias, amount=float(-delta) + ) + + def _sanitize_for_log(value: Any) -> str: """Strip CR/LF from user-controlled values to prevent log injection.""" try: @@ -2481,6 +2508,17 @@ async def team_member_add( raise HTTPException( status_code=404, detail={"error": f"Team with id {data.team_id} not found"} ) + + # Emit team-members gauge delta (+N) for the just-added members. Duplicates + # are rejected earlier by ``team_member_add_duplication_check``, so every + # entry in ``data.member`` corresponds to a real net-new membership. + members_added = 1 if isinstance(data.member, Member) else len(data.member) + _emit_team_members_metric_delta( + team_id=updated_team.team_id, + team_alias=updated_team.team_alias, + delta=members_added, + ) + return TeamAddMemberResponse( **updated_team.model_dump(), updated_users=updated_users, @@ -2679,6 +2717,15 @@ async def team_member_delete( } ) + # Emit team-members gauge delta (-1) for the removed member. The membership + # was verified to exist via ``is_member_in_team`` above, so this is always a + # real net deletion. + _emit_team_members_metric_delta( + team_id=existing_team_row.team_id, + team_alias=existing_team_row.team_alias, + delta=-1, + ) + return existing_team_row diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 43a287f29bc..317060de5ef 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -249,6 +249,8 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_check_batch_cost_jobs_processed_total", "litellm_check_batch_cost_errors_total", "litellm_check_batch_cost_last_run_timestamp", + # Team member metrics + "litellm_team_members_metric", ] @@ -701,6 +703,14 @@ class PrometheusMetricLabels: litellm_check_batch_cost_last_run_timestamp: List[str] = [] + # Team member count metric — incremented when a team member is added, + # decremented when one is removed. Labelled by team_id (``team``) and + # ``team_alias`` per the dashboards' team-level grouping. + litellm_team_members_metric = [ + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + ] + @staticmethod def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]: default_labels = getattr(PrometheusMetricLabels, label_name) 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 19ae819c85a..7af7cd2dcce 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -824,3 +824,85 @@ def test_custom_latency_buckets(): REGISTRY.unregister(collector) except Exception: pass + + +class TestPrometheusTeamMembersMetric: + """Tests for the per-team Prometheus team-members gauge. + + The gauge increments when a member is added to a team and decrements when + one is removed. It is labelled by ``team`` (team_id) and ``team_alias``. + """ + + def _gauge_value(self, prometheus_logger, team_id: str, team_alias: str) -> float: + """Read the current value of the team-members gauge for a label combo.""" + for metric in REGISTRY.collect(): + for sample in metric.samples: + if sample.name != "litellm_team_members_metric": + continue + if ( + sample.labels.get("team") == team_id + and sample.labels.get("team_alias") == team_alias + ): + return sample.value + return 0.0 + + def test_metric_is_initialized_with_team_labels(self, prometheus_logger): + """The gauge should exist and carry the two team-level labels.""" + assert hasattr(prometheus_logger, "litellm_team_members_metric") + assert prometheus_logger.litellm_team_members_metric is not None + # Sanity-check the label names by labelling and reading back the sample. + prometheus_logger.litellm_team_members_metric.labels( + team="t-init", team_alias="t-init-alias" + ).set(0) + for metric in REGISTRY.collect(): + for sample in metric.samples: + if ( + sample.name == "litellm_team_members_metric" + and sample.labels.get("team") == "t-init" + ): + assert set(sample.labels.keys()) >= {"team", "team_alias"} + return + pytest.fail("litellm_team_members_metric sample not found in registry") + + def test_increment_then_decrement_round_trips_to_zero(self, prometheus_logger): + """Add 3 members, then remove 3 — gauge should land back at 0.""" + team_id, team_alias = "t-rt", "rt-alias" + for _ in range(3): + prometheus_logger.increment_team_members_metric( + team_id=team_id, team_alias=team_alias + ) + assert self._gauge_value(prometheus_logger, team_id, team_alias) == 3.0 + for _ in range(3): + prometheus_logger.decrement_team_members_metric( + team_id=team_id, team_alias=team_alias + ) + assert self._gauge_value(prometheus_logger, team_id, team_alias) == 0.0 + + def test_increment_with_amount_supports_bulk_add(self, prometheus_logger): + """A single bulk add of N members should produce a +N delta.""" + prometheus_logger.increment_team_members_metric( + team_id="t-bulk", team_alias="bulk-alias", amount=5 + ) + assert self._gauge_value(prometheus_logger, "t-bulk", "bulk-alias") == 5.0 + + def test_independent_teams_track_independently(self, prometheus_logger): + """Counts on different teams must not bleed into each other.""" + prometheus_logger.increment_team_members_metric( + team_id="team-a", team_alias="alias-a" + ) + prometheus_logger.increment_team_members_metric( + team_id="team-a", team_alias="alias-a" + ) + prometheus_logger.increment_team_members_metric( + team_id="team-b", team_alias="alias-b" + ) + assert self._gauge_value(prometheus_logger, "team-a", "alias-a") == 2.0 + assert self._gauge_value(prometheus_logger, "team-b", "alias-b") == 1.0 + + def test_missing_team_alias_falls_back_to_empty_string(self, prometheus_logger): + """Teams without an alias must not crash the labels() call.""" + prometheus_logger.increment_team_members_metric( + team_id="t-noalias", team_alias=None + ) + # Empty string is the documented fallback when alias is None. + assert self._gauge_value(prometheus_logger, "t-noalias", "") == 1.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 5c7bbc46c95..f9efbef10bd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -7943,3 +7943,83 @@ async def test_team_member_me_returns_404_for_unknown_team(mock_db_client): user_api_key_dict=caller_auth, ) assert exc_info.value.status_code == 404 + + +class TestEmitTeamMembersMetricDelta: + """Tests for the ``_emit_team_members_metric_delta`` helper. + + The helper is the seam between the team management endpoints and the + Prometheus team-members gauge — it should call the right + ``PrometheusLogger`` method (inc or dec) based on sign, and silently + no-op when Prometheus is not configured. + """ + + def test_no_op_when_prometheus_logger_not_registered(self): + """Without a PrometheusLogger in callbacks the helper must not raise.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric_delta, + ) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.PrometheusLogger.get_instance", + return_value=None, + ): + # Should not raise, and there's nothing to assert on the call — the + # absence of an exception is the contract being tested. + _emit_team_members_metric_delta( + team_id="t-1", team_alias="alias-1", delta=3 + ) + + def test_zero_delta_short_circuits_before_lookup(self): + """delta=0 must not even touch get_instance — no work to do.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric_delta, + ) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.PrometheusLogger.get_instance" + ) as mock_get: + _emit_team_members_metric_delta( + team_id="t-1", team_alias="alias-1", delta=0 + ) + mock_get.assert_not_called() + + def test_positive_delta_calls_increment(self): + """A +N delta routes to ``increment_team_members_metric(amount=N)``.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric_delta, + ) + + mock_logger = MagicMock() + with patch( + "litellm.proxy.management_endpoints.team_endpoints.PrometheusLogger.get_instance", + return_value=mock_logger, + ): + _emit_team_members_metric_delta( + team_id="t-1", team_alias="alias-1", delta=3 + ) + + mock_logger.increment_team_members_metric.assert_called_once_with( + team_id="t-1", team_alias="alias-1", amount=3.0 + ) + mock_logger.decrement_team_members_metric.assert_not_called() + + def test_negative_delta_calls_decrement_with_absolute_amount(self): + """A -N delta routes to ``decrement_team_members_metric(amount=N)``.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _emit_team_members_metric_delta, + ) + + mock_logger = MagicMock() + with patch( + "litellm.proxy.management_endpoints.team_endpoints.PrometheusLogger.get_instance", + return_value=mock_logger, + ): + _emit_team_members_metric_delta( + team_id="t-1", team_alias="alias-1", delta=-1 + ) + + mock_logger.decrement_team_members_metric.assert_called_once_with( + team_id="t-1", team_alias="alias-1", amount=1.0 + ) + mock_logger.increment_team_members_metric.assert_not_called()