This commit is contained in:
Syed Ahnaf Raza 2026-08-26 21:51:01 +05:30 committed by GitHub
commit a5bcb44800
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 57 additions and 0 deletions

View file

@ -3672,6 +3672,8 @@ class PrometheusLogger(CustomLogger):
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
if spend is None and team_info.spend is not None:
team_object.spend = team_info.spend + response_cost
return team_object

View file

@ -359,6 +359,61 @@ async def test_assemble_team_object_does_not_override_metadata_max_budget(
), "max_budget from metadata must not be replaced by the DB value"
async def test_assemble_team_object_uses_db_spend_when_metadata_is_none(
prometheus_logger,
):
"""
When spend is None in request metadata (cached team object not hydrated),
_assemble_team_object must backfill it from get_team_object, otherwise the
remaining-budget gauge reports the full max_budget as if nothing was spent.
"""
db_team = MagicMock()
db_team.max_budget = 1000.0
db_team.spend = 902.60
db_team.budget_reset_at = None
with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: # test-quality-ok: stubbing the DB fetch, same pattern as the sibling max_budget fallback tests above
mock_get_team.return_value = db_team
team_object = await prometheus_logger._assemble_team_object(
team_id="team-1",
team_alias="team-a",
spend=None, # simulates None coming from request metadata
max_budget=1000.0,
response_cost=0.5,
)
assert (
team_object.spend == 902.60 + 0.5
), "spend should be populated from DB (plus request cost) when metadata value is None"
async def test_assemble_team_object_does_not_override_metadata_spend(
prometheus_logger,
):
"""
When spend IS present in request metadata, it must not be replaced by the
DB value.
"""
db_team = MagicMock()
db_team.max_budget = 1000.0
db_team.spend = 9999.0
db_team.budget_reset_at = None
with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: # test-quality-ok: stubbing the DB fetch, same pattern as the sibling max_budget fallback tests above
mock_get_team.return_value = db_team
team_object = await prometheus_logger._assemble_team_object(
team_id="team-1",
team_alias="team-a",
spend=50.0, # metadata has a real value
max_budget=1000.0,
response_cost=1.0,
)
assert (
team_object.spend == 50.0 + 1.0
), "spend 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,
):