scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

feat(prometheus): add litellm_team_member_count_metric gauge

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
oss-agent-shin 2026-05-19 14:27:23 -07:00
parent a72414a061
commit e438fee45e
4 changed files with 413 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
litellm/integrations/prometheus.py
View file

@ -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
########################################

24
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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 = {}

9
litellm/types/integrations/prometheus.py
View file

@ -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,

318
tests/test_litellm/integrations/test_prometheus_team_member_count_metric.py Normal file
View file

@ -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