refactored location of budget membership fix

Author: hagan

litellm/proxy/management_endpoints/common_utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Union
+from typing import Any, Union, Optional
 
 from litellm.proxy._types import (
     GenerateKeyRequest,
@@ -46,3 +46,77 @@ def _set_object_metadata_field(
         _premium_user_check()
     object_data.metadata = object_data.metadata or {}
     object_data.metadata[field_name] = value
+
+
+
+async def _upsert_budget_and_membership(
+    tx,
+    *,
+    team_id: str,
+    user_id: str,
+    max_budget: Optional[float],
+    existing_budget_id: Optional[str],
+    user_api_key_dict: UserAPIKeyAuth,
+):
+    """
+    Helper function to Create/Update or Delete the budget within the team membership
+    Args:
+        tx: The transaction object
+        team_id: The ID of the team
+        user_id: The ID of the user
+        max_budget: The maximum budget for the team
+        existing_budget_id: The ID of the existing budget, if any
+        user_api_key_dict: User API Key dictionary containing user information
+
+    If max_budget is None, the user's budget is removed from the team membership.
+    If max_budget exists, a budget is updated or created and linked to the team membership.
+    """
+    if max_budget is None:
+        # disconnect the budget since max_budget is None
+        await tx.litellm_teammembership.update(
+            where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
+            data={"litellm_budget_table": {"disconnect": True}},
+        )
+        return
+
+    if existing_budget_id:
+        # update the existing budget
+        await tx.litellm_budgettable.update(
+            where={"budget_id": existing_budget_id},
+            data={"max_budget": max_budget},
+        )
+        return
+
+    # create a new budget
+    new_budget = await tx.litellm_budgettable.create(
+        data={
+            "max_budget": max_budget,
+            "created_by": user_api_key_dict.user_id or "",
+            "updated_by": user_api_key_dict.user_id or "",
+        },
+        include={"team_membership": True},
+    )
+    # upsert the team membership with the new/updated budget
+    await tx.litellm_teammembership.upsert(
+        where={
+            "user_id_team_id": {
+                "user_id": user_id,
+                "team_id": team_id,
+            }
+        },
+        data={
+            "create": {
+                "user_id": user_id,
+                "team_id": team_id,
+                "litellm_budget_table": {
+                    "connect": {"budget_id": new_budget.budget_id},
+                },
+            },
+            "update": {
+                "litellm_budget_table": {
+                    "connect": {"budget_id": new_budget.budget_id},
+                },
+            },
+        },
+    )
+

litellm/proxy/management_endpoints/team_endpoints.py

@@ -66,6 +66,7 @@
     _is_user_team_admin,
     _set_object_metadata_field,
     _user_has_admin_view,
+    _upsert_budget_and_membership,
 )
 from litellm.proxy.management_endpoints.tag_management_endpoints import (
     get_daily_activity,
@@ -1266,56 +1267,13 @@ async def team_member_update(
 
     ### upsert new budget
     async with prisma_client.db.tx() as tx:
-        if data.max_budget_in_team is not None:
-            if identified_budget_id is None:
-                new_budget = await tx.litellm_budgettable.create(
-                    data={
-                        "max_budget": data.max_budget_in_team,
-                        "created_by": user_api_key_dict.user_id or "",
-                        "updated_by": user_api_key_dict.user_id or "",
-                    },
-                    include={"team_membership": True},
-                )
-                await tx.litellm_teammembership.upsert(
-                    where={
-                        "user_id_team_id": {
-                            "user_id": received_user_id,
-                            "team_id": data.team_id,
-                        }
-                    },
-                    data={
-                        "create": {
-                            "user_id": received_user_id,
-                            "team_id": data.team_id,
-                            "litellm_budget_table": {
-                                "connect": {"budget_id": new_budget.budget_id},
-                            },
-                        },
-                        "update": {
-                            "litellm_budget_table": {
-                                "connect": {"budget_id": new_budget.budget_id},
-                            },
-                        },
-                    },
-                )
-            elif identified_budget_id is not None:
-                await prisma_client.db.litellm_budgettable.update(
-                    where={"budget_id": identified_budget_id},
-                    data={"max_budget": data.max_budget_in_team},
-                )
-        else:
-            await tx.litellm_teammembership.update(
-            where={
-                "user_id_team_id": {
-                    "user_id": received_user_id,
-                    "team_id": data.team_id,
-                }
-            },
-            data={
-                "litellm_budget_table": {
-                    "disconnect": True
-                }
-            },
+        await _upsert_budget_and_membership(
+            tx=tx,
+            team_id=data.team_id,
+            user_id=received_user_id,
+            max_budget=data.max_budget_in_team,
+            existing_budget_id=identified_budget_id,
+            user_api_key_dict=user_api_key_dict,
         )
 
     ### update team member role