From e7d537442d5aeac01d9085b19bddd1011ae2ecb7 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 20:13:52 +0000 Subject: [PATCH] fix(prometheus): read the cached customer row for request-time budget gauges The request path used get_end_user_object, which falls back to a database lookup on a cache miss. Read the LiteLLM_EndUserTable row auth already cached instead, with the default budget already attached, and leave misses to the periodic refresh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/prometheus.py | 21 ++-- .../test_prometheus_user_team_metrics.py | 111 ++++++++++++++---- 2 files changed, 100 insertions(+), 32 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 4f67cb6b0e2..bb38b25ab99 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, ) +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import ( LiteLLM_DeletedVerificationToken, LiteLLM_TeamTable, @@ -4209,29 +4210,25 @@ class PrometheusLogger(CustomLogger): if not end_user_id: return - from litellm.proxy.auth.auth_checks import get_end_user_object - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - if prisma_client is None: - return + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + from litellm.proxy.proxy_server import user_api_key_cache try: - end_user_object: Final = await get_end_user_object( - end_user_id=end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + cached_customer: Final = await user_api_key_cache.async_get_cache( + key=end_user_cache_key(end_user_id), + model_type=LiteLLM_EndUserTable, ) except Exception as e: verbose_logger.debug("[Non-Blocking] Prometheus: Error getting customer info: %s", e) return - if end_user_object is None: + if cached_customer is None: return - budget_table: Final = end_user_object.litellm_budget_table + budget_table: Final = cached_customer.litellm_budget_table self._set_customer_budget_metrics( end_user_id=end_user_id, - spend=end_user_object.spend + response_cost, + spend=cached_customer.spend + response_cost, max_budget=budget_table.max_budget if budget_table is not None else None, budget_reset_at=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 077ea305053..7969c4741ef 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -1015,7 +1015,7 @@ def test_set_customer_budget_metrics_without_budget_only_emits_remaining(prometh @pytest.mark.asyncio -async def test_increment_remaining_budget_metrics_emits_customer_gauges_from_end_user_object( +async def test_increment_remaining_budget_metrics_emits_customer_gauges_from_cached_end_user( prometheus_logger, customer_metrics_enabled ): import sys @@ -1030,14 +1030,14 @@ async def test_increment_remaining_budget_metrics_emits_customer_gauges_from_end budget_id="budget-1", litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget-1", max_budget=1000.0), ) - get_end_user_object = AsyncMock(return_value=end_user) + get_end_user_object = AsyncMock() mock_proxy_server = MagicMock() - mock_proxy_server.prisma_client = MagicMock() - mock_proxy_server.user_api_key_cache = MagicMock() + mock_proxy_server.prisma_client = None + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user) with ( patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}), - patch("litellm.proxy.auth.auth_checks.get_end_user_object", get_end_user_object), # test-quality-ok: [TQ008] logger resolves customers through the proxy auth lookup, no injection seam + patch("litellm.proxy.auth.auth_checks.get_end_user_object", get_end_user_object), # test-quality-ok: [TQ008] assert the request path never reaches the DB-backed auth lookup ): await prometheus_logger._increment_remaining_budget_metrics( user_api_team=None, @@ -1049,30 +1049,104 @@ async def test_increment_remaining_budget_metrics_emits_customer_gauges_from_end end_user_id="cust-req", ) - get_end_user_object.assert_awaited_once() - assert get_end_user_object.await_args.kwargs["end_user_id"] == "cust-req" + get_end_user_object.assert_not_awaited() + cache_read = mock_proxy_server.user_api_key_cache.async_get_cache + cache_read.assert_awaited_once() + assert cache_read.await_args.kwargs["key"] == "end_user_id:cust-req" assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-req") == pytest.approx(650.0) assert _customer_sample("litellm_customer_max_budget_metric", "cust-req") == pytest.approx(1000.0) +@pytest.mark.asyncio +async def test_set_customer_budget_metrics_after_api_request_uses_cached_default_budget( + prometheus_logger, customer_metrics_enabled +): + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + + end_user = LiteLLM_EndUserTable( + user_id="cust-default", + blocked=False, + spend=0.5, + budget_id=None, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="default-budget", max_budget=3.0), + ) + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-default", + response_cost=0.5, + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-default") == pytest.approx(2.0) + assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0) + + +@pytest.mark.asyncio +async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining( + prometheus_logger, customer_metrics_enabled +): + import sys + + from litellm.models.end_user import LiteLLM_EndUserTable + + end_user = LiteLLM_EndUserTable(user_id="cust-no-budget", blocked=False, spend=2.0, budget_id=None) + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-no-budget", + response_cost=1.0, + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-no-budget") == float("inf") + assert _customer_sample("litellm_customer_max_budget_metric", "cust-no-budget") is None + + +@pytest.mark.asyncio +async def test_set_customer_budget_metrics_after_api_request_skips_uncached_customer( + prometheus_logger, customer_metrics_enabled +): + import sys + + get_end_user_object = AsyncMock() + mock_proxy_server = MagicMock() + mock_proxy_server.prisma_client = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=None) + + with ( + patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}), + patch("litellm.proxy.auth.auth_checks.get_end_user_object", get_end_user_object), # test-quality-ok: [TQ008] assert a cache miss does not fall back to the DB-backed auth lookup + ): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-uncached", + response_cost=1.0, + ) + + get_end_user_object.assert_not_awaited() + mock_proxy_server.prisma_client.assert_not_called() + assert prometheus_logger.litellm_remaining_customer_budget_metric._metrics == {} + + @pytest.mark.asyncio async def test_set_customer_budget_metrics_after_api_request_without_end_user_is_noop(prometheus_logger): import sys - get_end_user_object = AsyncMock() mock_proxy_server = MagicMock() - mock_proxy_server.prisma_client = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock() - with ( - patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}), - patch("litellm.proxy.auth.auth_checks.get_end_user_object", get_end_user_object), # test-quality-ok: [TQ008] assert the proxy auth lookup is never reached without an end user - ): + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): await prometheus_logger._set_customer_budget_metrics_after_api_request( end_user_id=None, response_cost=1.0, ) - get_end_user_object.assert_not_awaited() + mock_proxy_server.user_api_key_cache.async_get_cache.assert_not_awaited() assert prometheus_logger.litellm_remaining_customer_budget_metric._metrics == {} @@ -1241,13 +1315,10 @@ async def test_customer_max_budget_gauge_emitted_when_only_it_is_configured(cust litellm_budget_table=LiteLLM_BudgetTable(budget_id="budget-1", max_budget=40.0), ) mock_proxy_server = MagicMock() - mock_proxy_server.prisma_client = MagicMock() - mock_proxy_server.user_api_key_cache = MagicMock() + mock_proxy_server.prisma_client = None + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock(return_value=end_user) - with ( - patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}), - patch("litellm.proxy.auth.auth_checks.get_end_user_object", AsyncMock(return_value=end_user)), # test-quality-ok: [TQ008] logger resolves customers through the proxy auth lookup, no injection seam - ): + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): await logger._increment_remaining_budget_metrics( user_api_team=None, user_api_team_alias=None,