From ca04d9915263b56b0c6c6b76d270bd36cf3dd9b6 Mon Sep 17 00:00:00 2001 From: amarrtech <272048731+amarrtech@users.noreply.github.com> Date: Fri, 25 Sep 2026 08:21:05 -0700 Subject: [PATCH 1/3] fix(prometheus): include unscoped keys in budget metrics Signed-off-by: amarrtech <272048731+amarrtech@users.noreply.github.com> --- litellm/integrations/prometheus.py | 2 -- .../test_prometheus_user_team_metrics.py | 20 +++++++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index c7bf291a887..4192edf42bc 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3684,7 +3684,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, ) @@ -3708,7 +3707,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/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 0fc91748af2..680253a904b 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/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 @@ -191,6 +193,24 @@ 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): + mock_find_many: Final = AsyncMock(return_value=[]) + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = mock_find_many + mock_prisma.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + + with patch("litellm.proxy.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 + 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") From e418466feda99edefc90dcbfc72330634bdfbc02 Mon Sep 17 00:00:00 2001 From: amarrtech <272048731+amarrtech@users.noreply.github.com> Date: Fri, 25 Sep 2026 08:36:48 -0700 Subject: [PATCH 2/3] test(prometheus): verify unscoped key budget gauges Signed-off-by: amarrtech <272048731+amarrtech@users.noreply.github.com> --- .../test_prometheus_user_team_metrics.py | 30 +++++++++++++++++-- 1 file changed, 27 insertions(+), 3 deletions(-) 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 680253a904b..a7aa8fb8a7d 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -195,12 +195,24 @@ class TestPrometheusUserTeamCountMetrics: @pytest.mark.asyncio async def test_initialize_api_key_budget_metrics_keeps_keys_without_team(self, prometheus_logger): - mock_find_many: Final = AsyncMock(return_value=[]) + 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=0) + mock_prisma.db.litellm_verificationtoken.count = AsyncMock(return_value=1) + prometheus_logger.litellm_remaining_api_key_budget_metric = MagicMock() + prometheus_logger.litellm_api_key_max_budget_metric = MagicMock() - with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(proxy_server, "prisma_client", mock_prisma): await prometheus_logger._initialize_api_key_budget_metrics() expected_where: Final = { @@ -210,6 +222,18 @@ class TestPrometheusUserTeamCountMetrics: ] } assert mock_find_many.await_args.kwargs["where"] == expected_where + prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once() + ( + prometheus_logger.litellm_remaining_api_key_budget_metric.labels.return_value.set.assert_called_once_with( + 75.0 + ) + ) + prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once() + ( + prometheus_logger.litellm_api_key_max_budget_metric.labels.return_value.set.assert_called_once_with( + 100.0 + ) + ) def test_active_users_metric_initialized(self, prometheus_logger): """litellm_active_users gauge must exist alongside litellm_total_users.""" From 11a0012d88e4f5fe3560c26bba9e8d9bb1f54433 Mon Sep 17 00:00:00 2001 From: Amrinder Randhawa <272048731+amarrtech@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:12:50 -0700 Subject: [PATCH 3/3] test(prometheus): assert emitted key budget gauges Signed-off-by: Amrinder Randhawa <272048731+amarrtech@users.noreply.github.com> --- .../test_prometheus_user_team_metrics.py | 20 ++++++------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index fe1e0fd0bbb..ca6929260d1 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -188,8 +188,6 @@ class TestPrometheusUserTeamCountMetrics: mock_prisma: Final = MagicMock() mock_prisma.db.litellm_verificationtoken.find_many = mock_find_many mock_prisma.db.litellm_verificationtoken.count = AsyncMock(return_value=1) - prometheus_logger.litellm_remaining_api_key_budget_metric = MagicMock() - prometheus_logger.litellm_api_key_max_budget_metric = MagicMock() with patch.object(proxy_server, "prisma_client", mock_prisma): await prometheus_logger._initialize_api_key_budget_metrics() @@ -201,18 +199,12 @@ class TestPrometheusUserTeamCountMetrics: ] } assert mock_find_many.await_args.kwargs["where"] == expected_where - prometheus_logger.litellm_remaining_api_key_budget_metric.labels.assert_called_once() - ( - prometheus_logger.litellm_remaining_api_key_budget_metric.labels.return_value.set.assert_called_once_with( - 75.0 - ) - ) - prometheus_logger.litellm_api_key_max_budget_metric.labels.assert_called_once() - ( - prometheus_logger.litellm_api_key_max_budget_metric.labels.return_value.set.assert_called_once_with( - 100.0 - ) - ) + 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."""