diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index c65c344992d..33bfa1f8b44 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -591,12 +591,22 @@ async def _upsert_budget_and_membership( "updated_by": user_api_key_dict.user_id or "", } - if is_shared_default: - default_budget_row: Final = await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) - if default_budget_row is not None: - default_budget_dict: Final = default_budget_row.model_dump() + seed_row_id: Final = ( + existing_budget_id + if is_shared_default + else team_default_budget_id + if team_default_budget_id is not None + and ("temp_budget_increase" in write_data or "temp_budget_expiry" in write_data) + else None + ) + if seed_row_id is not None: + seed_row: Final = await tx.litellm_budgettable.find_unique(where={"budget_id": seed_row_id}) + if seed_row is not None: + seed_dict: Final = seed_row.model_dump() for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: - value = default_budget_dict.get(field) + value = seed_dict.get(field) + if field == "max_budget" and value == 0 and not is_shared_default: + continue if _is_set_budget_value(value): create_data[field] = value