diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 19d161b8aa4..8ce94af55e5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3995,7 +3995,6 @@ class PrometheusLogger(CustomLogger): """ Initialize API key budget metrics by reusing the generic pagination logic. """ - from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.management_endpoints.key_management_endpoints import ( list_key_helper, ) @@ -4019,7 +4018,6 @@ class PrometheusLogger(CustomLogger): team_id=None, key_alias=None, key_hash=None, - exclude_team_id=UI_SESSION_TOKEN_TEAM_ID, return_full_object=True, organization_id=None, ) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index ab0fb67b52f..ca6929260d1 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -3,11 +3,13 @@ Unit tests for Prometheus user and team count metrics """ from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.integrations.prometheus import PrometheusLogger @@ -170,6 +172,40 @@ class TestPrometheusUserTeamCountMetrics: prometheus_logger._initialize_api_key_budget_metrics.assert_called_once() prometheus_logger._initialize_user_and_team_count_metrics.assert_called_once() + @pytest.mark.asyncio + async def test_initialize_api_key_budget_metrics_keeps_keys_without_team(self, prometheus_logger): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + unscoped_key: Final = UserAPIKeyAuth( + token="hashed-unscoped-key", + key_alias="unscoped-key", + team_id=None, + spend=25.0, + max_budget=100.0, + ) + mock_find_many: Final = AsyncMock(return_value=[unscoped_key]) + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = mock_find_many + mock_prisma.db.litellm_verificationtoken.count = AsyncMock(return_value=1) + + with patch.object(proxy_server, "prisma_client", mock_prisma): + await prometheus_logger._initialize_api_key_budget_metrics() + + expected_where: Final = { + "OR": [ + {"team_id": None}, + {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, + ] + } + assert mock_find_many.await_args.kwargs["where"] == expected_where + labels: Final = { + "hashed_api_key": "hashed-unscoped-key", + "api_key_alias": "unscoped-key", + } + assert REGISTRY.get_sample_value("litellm_remaining_api_key_budget_metric", labels) == 75.0 + assert REGISTRY.get_sample_value("litellm_api_key_max_budget_metric", labels) == 100.0 + def test_active_users_metric_initialized(self, prometheus_logger): """litellm_active_users gauge must exist alongside litellm_total_users.""" assert hasattr(prometheus_logger, "litellm_active_users_metric")