diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index dc53a39ea99..50ff8cd6f32 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -4230,30 +4230,30 @@ class PrometheusLogger(CustomLogger): max_budget=max_budget if max_budget is not None else carried.max_budget, budget_reset_at=carried.budget_reset_at, ) - team_object: Final = LiteLLM_TeamTable( - team_id=team_id, - team_alias=team_alias, - spend=_total_team_spend, - max_budget=max_budget, - ) + team_info: LiteLLM_TeamTable | None = None try: - team_info: Final = await get_team_object( + team_info = await get_team_object( team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, ) except Exception as e: verbose_logger.debug("[Non-Blocking] Prometheus: Error getting team info: %s", e) - return team_object - 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 - if spend is None and team_info.spend is not None: - team_object.spend = team_info.spend + response_cost - - return team_object + # metadata spend of None means "unknown", not "nothing spent" -- fall back to the DB value + _base_spend: Final = ( + spend if spend is not None else (team_info.spend if team_info and team_info.spend is not None else 0) + ) + _resolved_max_budget: Final = ( + max_budget if max_budget is not None else (team_info.max_budget if team_info else None) + ) + return LiteLLM_TeamTable( + team_id=team_id, + team_alias=team_alias, + spend=_base_spend + response_cost, + max_budget=_resolved_max_budget, + budget_reset_at=team_info.budget_reset_at if team_info else None, + ) def _set_team_budget_metrics( self, diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index e10caae1c96..3b9007f21e3 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -356,7 +356,7 @@ async def test_assemble_team_object_uses_db_spend_when_metadata_is_none( team_object = await prometheus_logger._assemble_team_object( team_id="team-1", team_alias="team-a", - spend=None, # simulates None coming from request metadata + spend=None, max_budget=1000.0, response_cost=0.5, ) @@ -383,7 +383,7 @@ async def test_assemble_team_object_does_not_override_metadata_spend( team_object = await prometheus_logger._assemble_team_object( team_id="team-1", team_alias="team-a", - spend=50.0, # metadata has a real value + spend=50.0, max_budget=1000.0, response_cost=1.0, )