From 6a4e29f87b333ae9977e8f878960e63becd89150 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 20 Jun 2025 17:57:20 -0700
Subject: [PATCH] fix(team_endpoints.py): support updating team budget on UI

---
 litellm/proxy/management_endpoints/team_endpoints.py   | 6 +++++-
 ui/litellm-dashboard/src/components/team/team_info.tsx | 8 ++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 14d04ede2c4..c6637d2aee5 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -193,7 +193,7 @@ async def _upsert_team_member_budget_table(
         team_table.metadata = {}
 
     team_member_budget_id = team_table.metadata.get("team_member_budget_id")
-    if team_member_budget_id is None and isinstance(team_member_budget_id, str):
+    if team_member_budget_id is not None and isinstance(team_member_budget_id, str):
         # Budget exists
         budget_row = await update_budget(
             budget_obj=BudgetNewRequest(
@@ -202,9 +202,13 @@ async def _upsert_team_member_budget_table(
             ),
             user_api_key_dict=user_api_key_dict,
         )
+        verbose_proxy_logger.info(
+            f"Updated team member budget table: {budget_row.budget_id}, with team_member_budget={team_member_budget}"
+        )
         if updated_kv.get("metadata") is None:
             updated_kv["metadata"] = {}
         updated_kv["metadata"]["team_member_budget_id"] = budget_row.budget_id
+        updated_kv.pop("team_member_budget", None)
     else:  # budget does not exist
         updated_kv = await _create_team_member_budget_table(
             data=team_table,
diff --git a/ui/litellm-dashboard/src/components/team/team_info.tsx b/ui/litellm-dashboard/src/components/team/team_info.tsx
index 1c6bfd46832..6f119e7d5c4 100644
--- a/ui/litellm-dashboard/src/components/team/team_info.tsx
+++ b/ui/litellm-dashboard/src/components/team/team_info.tsx
@@ -235,6 +235,10 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
         organization_id: values.organization_id,
       };
 
+      if (values.team_member_budget !== undefined) {
+        updateData.team_member_budget = Number(values.team_member_budget);
+      }
+
       // Handle object_permission updates
       if (values.vector_stores !== undefined || values.mcp_servers !== undefined) {
         updateData.object_permission = {
@@ -420,6 +424,10 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                     <NumericalInput step={0.01} precision={2} style={{ width: "100%" }} />
                   </Form.Item>
 
+                  <Form.Item label="Team Member Budget (USD)" name="team_member_budget">
+                    <NumericalInput step={0.01} precision={2} style={{ width: "100%" }} />
+                  </Form.Item>
+
                   <Form.Item label="Reset Budget" name="budget_duration">
                     <Select placeholder="n/a">
                       <Select.Option value="24h">daily</Select.Option>