From e438fee45e7ed835e5c0657117f8509a0fd86a77 Mon Sep 17 00:00:00 2001 From: oss-agent-shin Date: Tue, 19 May 2026 14:27:23 -0700 Subject: [PATCH] feat(prometheus): add litellm_team_member_count_metric gauge Adds a new Gauge metric that tracks the number of members per team. The gauge is incremented when a user is added via /team/member_add and decremented when one is removed via /team/member_delete. Labels: - team (team_id) - team_alias (human-readable alias) Implementation: - Defines the gauge in PrometheusLogger.__init__ using the existing label-aware factory. - Adds static helpers emit_team_member_added_metric / emit_team_member_removed_metric that no-op when no PrometheusLogger is registered (matches the rest of the codebase's pattern). - In _add_team_members_to_team, emits the actual delta (post-dedup) so duplicates filtered by _update_team_members_list don't double-count. - In team_member_delete, emits the actual delta after the team row update succeeds. Tests (16 cases) cover gauge definition, label set, inc/dec helpers, no-op guards (missing logger, missing team_id, non-positive count) and endpoint-level integration showing the gauge ending at the expected value after adds and a delete. Co-Authored-By: Claude Opus 4.7 --- litellm/integrations/prometheus.py | 62 ++++ .../management_endpoints/team_endpoints.py | 24 ++ litellm/types/integrations/prometheus.py | 9 + ...est_prometheus_team_member_count_metric.py | 318 ++++++++++++++++++ 4 files changed, 413 insertions(+) create mode 100644 tests/test_litellm/integrations/test_prometheus_team_member_count_metric.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 30af0dcb8ed..2f351838428 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -68,6 +68,58 @@ class PrometheusLogger(CustomLogger): return cb return None + @staticmethod + def emit_team_member_added_metric( + team_id: Optional[str], + team_alias: Optional[str], + count: int = 1, + ) -> None: + """ + Increment the team member count gauge when one or more members are added + to a team. No-op if PrometheusLogger is not registered as a callback or + if count <= 0 (e.g. all members were duplicates). + """ + if count <= 0 or not team_id: + return + instance = PrometheusLogger.get_instance() + if instance is None: + return + try: + instance.litellm_team_member_count_metric.labels( + team=team_id, + team_alias=team_alias or "", + ).inc(count) + except Exception as e: + verbose_logger.debug( + f"Prometheus: error emitting team member added metric: {e}" + ) + + @staticmethod + def emit_team_member_removed_metric( + team_id: Optional[str], + team_alias: Optional[str], + count: int = 1, + ) -> None: + """ + Decrement the team member count gauge when one or more members are + removed from a team. No-op if PrometheusLogger is not registered as a + callback or if count <= 0. + """ + if count <= 0 or not team_id: + return + instance = PrometheusLogger.get_instance() + if instance is None: + return + try: + instance.litellm_team_member_count_metric.labels( + team=team_id, + team_alias=team_alias or "", + ).dec(count) + except Exception as e: + verbose_logger.debug( + f"Prometheus: error emitting team member removed metric: {e}" + ) + def __init__( # noqa: PLR0915 self, **kwargs, @@ -477,6 +529,16 @@ class PrometheusLogger(CustomLogger): labelnames=[], ) + # Team membership gauge - tracks number of members per team. + # Incremented on /team/member_add and decremented on /team/member_delete. + self.litellm_team_member_count_metric = self._gauge_factory( + "litellm_team_member_count_metric", + "Number of members in a team. Incremented when a member is added, decremented when a member is removed.", + labelnames=self.get_labels_for_metric( + "litellm_team_member_count_metric" + ), + ) + ######################################## # Managed Batch Metrics ######################################## diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 35e3d196e9e..2b76bd429aa 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, @@ -2237,6 +2238,10 @@ async def _add_team_members_to_team( litellm_proxy_admin_name: str, ) -> Tuple[LiteLLM_TeamTable, List[LiteLLM_UserTable], List[LiteLLM_TeamMembership]]: """Add team members to the team.""" + # Track size before adding so we can emit the actual delta (ignoring + # duplicates that get filtered out by _update_team_members_list). + members_before = len(complete_team_data.members_with_roles) + # Process and add new members updated_users, updated_team_memberships = await _process_team_members( data=data, @@ -2260,6 +2265,15 @@ async def _add_team_members_to_team( data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) + # Emit prometheus team member count metric for the actual delta. + members_added = len(complete_team_data.members_with_roles) - members_before + if members_added > 0: + PrometheusLogger.emit_team_member_added_metric( + team_id=complete_team_data.team_id, + team_alias=complete_team_data.team_alias, + count=members_added, + ) + return updated_team, updated_users, updated_team_memberships @@ -2589,6 +2603,7 @@ async def team_member_delete( ) ## DELETE MEMBER FROM TEAM + members_before = len(existing_team_row.members_with_roles) is_member_in_team, new_team_members = _cleanup_members_with_roles( existing_team_row=existing_team_row, data=data, @@ -2608,6 +2623,15 @@ async def team_member_delete( data={"members_with_roles": json.dumps(_db_new_team_members)}, # type: ignore ) + # Emit prometheus team member count metric for the actual delta. + members_removed = members_before - len(new_team_members) + if members_removed > 0: + PrometheusLogger.emit_team_member_removed_metric( + team_id=existing_team_row.team_id, + team_alias=existing_team_row.team_alias, + count=members_removed, + ) + ## 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 43a287f29bc..fd392d245c9 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -239,6 +239,8 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", "litellm_in_flight_requests", + # Team membership metrics + "litellm_team_member_count_metric", # Managed batch metrics "litellm_managed_batch_created_total", "litellm_managed_file_size_bytes", @@ -503,6 +505,13 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, ] + # Team membership count - emitted on team/member_add and team/member_delete. + # Gauge value increases when a user is added and decreases when one is removed. + litellm_team_member_count_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_team_member_count_metric.py b/tests/test_litellm/integrations/test_prometheus_team_member_count_metric.py new file mode 100644 index 00000000000..5029a84cfbd --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_team_member_count_metric.py @@ -0,0 +1,318 @@ +""" +Unit tests for the litellm_team_member_count_metric Prometheus gauge. + +The gauge tracks the number of members currently in each team: +- Incremented when a member is added via /team/member_add +- Decremented when a member is removed via /team/member_delete +- Labels: ``team`` (team_id) and ``team_alias`` +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.types.integrations.prometheus import PrometheusMetricLabels + + +@pytest.fixture(autouse=True) +def cleanup_prometheus_registry(): + """Clear the registry between tests so PrometheusLogger() can be re-created.""" + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + # Reset callbacks so emit helpers don't pick up loggers from prior tests + original_callbacks = list(litellm.callbacks) + litellm.callbacks = [] + yield + litellm.callbacks = original_callbacks + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +@pytest.fixture +def prometheus_logger(): + """Fresh PrometheusLogger registered as a litellm callback.""" + logger = PrometheusLogger() + litellm.callbacks = [logger] + return logger + + +def _get_gauge_value(team_id: str, team_alias: str) -> float: + """Read the current value of litellm_team_member_count_metric for a label set.""" + for metric in REGISTRY.collect(): + for sample in metric.samples: + if ( + sample.name == "litellm_team_member_count_metric" + and sample.labels.get("team") == team_id + and sample.labels.get("team_alias") == team_alias + ): + return sample.value + return 0.0 + + +class TestTeamMemberCountMetricDefinition: + """Verify the gauge is declared with the expected labels and shape.""" + + def test_metric_attribute_exists(self, prometheus_logger): + assert hasattr(prometheus_logger, "litellm_team_member_count_metric") + assert prometheus_logger.litellm_team_member_count_metric is not None + + def test_metric_labels_are_team_and_team_alias(self): + labels = PrometheusMetricLabels.litellm_team_member_count_metric + assert "team" in labels + assert "team_alias" in labels + assert set(labels) == {"team", "team_alias"} + + def test_metric_is_a_gauge(self, prometheus_logger): + """Gauge supports both .inc and .dec — Counter would not.""" + metric = prometheus_logger.litellm_team_member_count_metric + labelled = metric.labels(team="t", team_alias="a") + assert hasattr(labelled, "inc") + assert hasattr(labelled, "dec") + + def test_metric_in_defined_prometheus_metrics_literal(self): + """The metric name must be in DEFINED_PROMETHEUS_METRICS so + get_labels_for_metric resolves it.""" + from typing import get_args + + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + assert "litellm_team_member_count_metric" in get_args( + DEFINED_PROMETHEUS_METRICS + ) + + +class TestEmitHelpers: + """The static emit_* helpers are the entry point used from endpoint code.""" + + def test_emit_added_increments_gauge(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=1 + ) + assert _get_gauge_value("team-abc", "my-team") == 1.0 + + def test_emit_added_supports_count_greater_than_one(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=3 + ) + assert _get_gauge_value("team-abc", "my-team") == 3.0 + + def test_emit_removed_decrements_gauge(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=5 + ) + PrometheusLogger.emit_team_member_removed_metric( + team_id="team-abc", team_alias="my-team", count=2 + ) + assert _get_gauge_value("team-abc", "my-team") == 3.0 + + def test_emit_add_then_remove_to_zero(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=2 + ) + PrometheusLogger.emit_team_member_removed_metric( + team_id="team-abc", team_alias="my-team", count=2 + ) + assert _get_gauge_value("team-abc", "my-team") == 0.0 + + def test_emit_zero_count_is_noop(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=0 + ) + PrometheusLogger.emit_team_member_removed_metric( + team_id="team-abc", team_alias="my-team", count=0 + ) + assert _get_gauge_value("team-abc", "my-team") == 0.0 + + def test_emit_negative_count_is_noop(self, prometheus_logger): + """Guards against accidental sign flips at the call site.""" + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=-3 + ) + assert _get_gauge_value("team-abc", "my-team") == 0.0 + + def test_emit_missing_team_id_is_noop(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id=None, team_alias="my-team", count=1 + ) + PrometheusLogger.emit_team_member_added_metric( + team_id="", team_alias="my-team", count=1 + ) + # Nothing should have been recorded under any team label + for metric in REGISTRY.collect(): + for sample in metric.samples: + assert sample.name != "litellm_team_member_count_metric" or ( + sample.value == 0.0 + ) + + def test_emit_none_team_alias_is_handled(self, prometheus_logger): + """team_alias is optional; we should fall back to an empty string label.""" + PrometheusLogger.emit_team_member_added_metric( + team_id="team-xyz", team_alias=None, count=1 + ) + assert _get_gauge_value("team-xyz", "") == 1.0 + + def test_emit_is_noop_when_logger_not_registered(self): + """If no PrometheusLogger is in litellm.callbacks the helpers must not raise.""" + litellm.callbacks = [] + # Must not raise and must not need a logger to exist + PrometheusLogger.emit_team_member_added_metric( + team_id="team-abc", team_alias="my-team", count=1 + ) + PrometheusLogger.emit_team_member_removed_metric( + team_id="team-abc", team_alias="my-team", count=1 + ) + + def test_per_team_labels_are_independent(self, prometheus_logger): + PrometheusLogger.emit_team_member_added_metric( + team_id="team-a", team_alias="alias-a", count=4 + ) + PrometheusLogger.emit_team_member_added_metric( + team_id="team-b", team_alias="alias-b", count=1 + ) + PrometheusLogger.emit_team_member_removed_metric( + team_id="team-a", team_alias="alias-a", count=1 + ) + assert _get_gauge_value("team-a", "alias-a") == 3.0 + assert _get_gauge_value("team-b", "alias-b") == 1.0 + + +# --------------------------------------------------------------------------- +# Endpoint integration: verify the helpers are invoked from the team endpoints. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_add_team_members_emits_metric_for_actual_delta(prometheus_logger): + """ + _add_team_members_to_team must emit the actual delta (post-dedup), not the + raw request size. We seed a team with 1 existing member, attempt to add 2 + members where 1 is a duplicate, and expect the gauge to increase by 1. + """ + from litellm.proxy._types import ( + LiteLLM_TeamTable, + Member, + TeamMemberAddRequest, + UserAPIKeyAuth, + LitellmUserRoles, + ) + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + team = LiteLLM_TeamTable( + team_id="team-add-1", + team_alias="add-team-alias", + members_with_roles=[Member(role="user", user_id="existing-user")], + ) + + data = TeamMemberAddRequest( + team_id="team-add-1", + member=[ + Member(role="user", user_id="existing-user"), # duplicate + Member(role="user", user_id="new-user"), # actually new + ], + ) + + prisma_client = MagicMock() + prisma_client.db = MagicMock() + prisma_client.db.litellm_teamtable = MagicMock() + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team) + + fake_user = MagicMock() + fake_user.user_id = "new-user" + fake_user.user_email = None + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._process_team_members", + new=AsyncMock(return_value=([fake_user], [])), + ), + ): + await _add_team_members_to_team( + data=data, + complete_team_data=team, + prisma_client=prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + # Only the new user counts toward the metric — the duplicate is filtered. + assert _get_gauge_value("team-add-1", "add-team-alias") == 1.0 + + +@pytest.mark.asyncio +async def test_team_member_delete_decrements_metric(prometheus_logger): + """ + team_member_delete must decrement the gauge by 1 for the removed member. + Pre-seed the gauge so we can verify the decrement against a non-zero start. + """ + from litellm.proxy._types import ( + LitellmUserRoles, + TeamMemberDeleteRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import team_member_delete + + team_id = "team-del-metric-1" + team_alias = "del-team-alias" + user_id = "user-to-remove@example.com" + + # Seed gauge at 3 — simulates a team with 3 members tracked. + PrometheusLogger.emit_team_member_added_metric( + team_id=team_id, team_alias=team_alias, count=3 + ) + + mock_team_row = MagicMock() + mock_team_row.model_dump.return_value = { + "team_id": team_id, + "team_alias": team_alias, + "members_with_roles": [ + {"user_id": user_id, "user_email": None, "role": "user"}, + {"user_id": "kept-user-1", "user_email": None, "role": "user"}, + {"user_id": "kept-user-2", "user_email": None, "role": "user"}, + ], + "team_member_permissions": [], + "metadata": {}, + "models": [], + "spend": 0.0, + } + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_teamtable = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team_row) + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row) + mock_prisma.db.litellm_usertable = MagicMock() + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_teammembership = MagicMock() + mock_prisma.db.litellm_teammembership.delete_many = AsyncMock( + return_value=MagicMock() + ) + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_verificationtoken.delete_many = AsyncMock( + return_value=MagicMock() + ) + + # team_member_delete imports prisma_client from litellm.proxy.proxy_server. + # Stub the module so we don't need to load the full proxy server. + fake_proxy_server = MagicMock() + fake_proxy_server.prisma_client = mock_prisma + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}): + await team_member_delete( + data=TeamMemberDeleteRequest(team_id=team_id, user_id=user_id), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + # Started at 3, removed 1 → expect 2 + assert _get_gauge_value(team_id, team_alias) == 2.0