From dfbd57a829cb3000fa8dd22609f6d5eb8a18ca5f Mon Sep 17 00:00:00 2001 From: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Date: Fri, 8 May 2026 22:37:33 +0000 Subject: [PATCH] avoid user lookup for Prometheus email labels Co-authored-by: ishaan-berri --- litellm/integrations/prometheus.py | 14 ++++++++++--- .../test_prometheus_user_team_metrics.py | 20 +++++++++---------- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index d044ad3375d..c28d873880b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3503,6 +3503,17 @@ class PrometheusLogger(CustomLogger): - Set user budget metrics """ if user_id: + if user_spend is not None and user_max_budget is not None: + self._set_user_budget_metrics( + LiteLLM_UserTable( + user_id=user_id, + user_email=user_email, + spend=user_spend + response_cost, + max_budget=user_max_budget, + ) + ) + return + user_object = await self._assemble_user_object( user_id=user_id, user_email=user_email, @@ -3556,9 +3567,6 @@ class PrometheusLogger(CustomLogger): if user_info: user_object.budget_reset_at = user_info.budget_reset_at - user_info_email = getattr(user_info, "user_email", None) - if user_object.user_email is None and isinstance(user_info_email, str): - user_object.user_email = user_info_email if user_object.max_budget is None and user_info.max_budget is not None: user_object.max_budget = user_info.max_budget 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 0b941a3340d..a94c5d65611 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -653,32 +653,30 @@ def test_user_budget_metrics_emit_user_email_label(prometheus_logger): ) -async def test_set_user_budget_metrics_after_api_request_uses_db_user_email( +async def test_set_user_budget_metrics_after_api_request_uses_metadata_without_lookup( prometheus_logger, ): - """The request path should populate user_email from user lookup when needed.""" + """Complete request metadata should emit budget labels without user lookup.""" prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() prometheus_logger.litellm_user_max_budget_metric = MagicMock() prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() - db_user = MagicMock() - db_user.max_budget = 500.0 - db_user.budget_reset_at = datetime(2099, 1, 1, tzinfo=timezone.utc) - db_user.user_email = "db-user@example.com" - with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user: - mock_get_user.return_value = db_user await prometheus_logger._set_user_budget_metrics_after_api_request( user_id="user-abc-123", - user_email=None, + user_email="metadata-user@example.com", user_spend=120.0, - user_max_budget=None, + user_max_budget=500.0, response_cost=0.5, ) + mock_get_user.assert_not_called() prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with( user="user-abc-123", - user_email="db-user@example.com", + user_email="metadata-user@example.com", + ) + prometheus_logger.litellm_remaining_user_budget_metric.labels().set.assert_called_once_with( + 379.5 )