diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 08db77e8571..121431d2114 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2686,6 +2686,8 @@ class PrometheusLogger(CustomLogger): if team_info: team_object.budget_reset_at = team_info.budget_reset_at + if team_object.max_budget is None and team_info.max_budget is not None: + team_object.max_budget = team_info.max_budget return team_object 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 a840e2fe162..40e33995c66 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -1,7 +1,8 @@ """ Unit tests for Prometheus user and team count metrics """ -from unittest.mock import MagicMock +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch import pytest from prometheus_client import REGISTRY @@ -258,3 +259,137 @@ class TestPrometheusUserTeamCountMetrics: assert True except Exception as e: pytest.fail(f"Metrics should handle large values: {e}") + + +# --------------------------------------------------------------------------- +# Regression tests: team budget showing +Inf when user_api_key_team_max_budget +# is None in request metadata but the team has a real budget in the DB. +# --------------------------------------------------------------------------- + + +async def test_assemble_team_object_uses_db_max_budget_when_metadata_is_none( + prometheus_logger, +): + """ + When max_budget is None in request metadata (e.g. stale key cache), + _assemble_team_object must fall back to the value returned by get_team_object + so that _safe_get_remaining_budget does not return +Inf. + """ + db_team = MagicMock() + db_team.max_budget = 3000.0 + db_team.budget_reset_at = datetime(2026, 3, 1, tzinfo=timezone.utc) + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + team_object = await prometheus_logger._assemble_team_object( + team_id="c5c33858-4379-4c90-8733-d9c58c312c10", + team_alias="ai-ml-local_dev", + spend=1617.02, + max_budget=None, # simulates None coming from request metadata + response_cost=0.5, + ) + + assert team_object.max_budget == 3000.0, ( + "max_budget should be populated from DB when metadata value is None" + ) + assert team_object.budget_reset_at == datetime(2026, 3, 1, tzinfo=timezone.utc) + + +async def test_assemble_team_object_does_not_override_metadata_max_budget( + prometheus_logger, +): + """ + When max_budget IS present in request metadata, it must not be overridden + by the DB value. + """ + db_team = MagicMock() + db_team.max_budget = 9999.0 + db_team.budget_reset_at = None + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + team_object = await prometheus_logger._assemble_team_object( + team_id="team-1", + team_alias="my-team", + spend=50.0, + max_budget=100.0, # metadata has a real value + response_cost=1.0, + ) + + assert team_object.max_budget == 100.0, ( + "max_budget from metadata must not be replaced by the DB value" + ) + + +async def test_set_team_budget_metrics_after_api_request_no_inf_when_metadata_budget_none( + prometheus_logger, +): + """ + End-to-end: when user_api_key_team_max_budget is None in request metadata + but the team has a real budget in the DB, the metric must NOT be set to +Inf. + """ + prometheus_logger.litellm_remaining_team_budget_metric = MagicMock() + prometheus_logger.litellm_team_max_budget_metric = MagicMock() + prometheus_logger.litellm_team_budget_remaining_hours_metric = MagicMock() + + db_team = MagicMock() + db_team.max_budget = 3000.0 + db_team.budget_reset_at = datetime(2026, 3, 1, tzinfo=timezone.utc) + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + await prometheus_logger._set_team_budget_metrics_after_api_request( + user_api_team="c5c33858-4379-4c90-8733-d9c58c312c10", + user_api_team_alias="ai-ml-local_dev", + team_spend=1617.02, + team_max_budget=None, # simulates stale key cache + response_cost=0.5, + ) + + set_call_args = ( + prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args + ) + assert set_call_args is not None, "remaining_team_budget_metric.labels().set was not called" + actual_value = set_call_args[0][0] + assert actual_value != float("inf"), ( + f"remaining_team_budget_metric must not be +Inf when team has a real budget; got {actual_value}" + ) + expected = 3000.0 - 1617.02 - 0.5 + assert abs(actual_value - expected) < 0.01, ( + f"Expected remaining budget ~{expected}, got {actual_value}" + ) + + +async def test_set_team_budget_metrics_after_api_request_inf_when_genuinely_no_budget( + prometheus_logger, +): + """ + When the team genuinely has no budget (max_budget=None in both metadata and + DB), +Inf is the correct value and must be preserved. + """ + prometheus_logger.litellm_remaining_team_budget_metric = MagicMock() + prometheus_logger.litellm_team_max_budget_metric = MagicMock() + prometheus_logger.litellm_team_budget_remaining_hours_metric = MagicMock() + + db_team = MagicMock() + db_team.max_budget = None + db_team.budget_reset_at = None + + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = db_team + await prometheus_logger._set_team_budget_metrics_after_api_request( + user_api_team="team-no-budget", + user_api_team_alias="no-budget-team", + team_spend=10.0, + team_max_budget=None, + response_cost=1.0, + ) + + set_call_args = ( + prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args + ) + assert set_call_args is not None + actual_value = set_call_args[0][0] + assert actual_value == float("inf"), ( + "remaining_team_budget_metric should be +Inf when team truly has no budget" + )