From 18b09eb59426d604a11dfa4cbb678dbd95aa77ba Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 13 Sep 2025 16:34:26 -0700
Subject: [PATCH] fix: org budget update fix (#14541)

---
 litellm/proxy/_types.py                       |  2 +-
 .../organization_endpoints.py                 | 48 ++++++++++++++++++-
 .../test_organization_endpoints.py            | 35 ++++++++++++++
 3 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 5a1952686bc..e7f1b3043e1 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1918,7 +1918,7 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase):
     model_config = ConfigDict(protected_namespaces=())
 
 
-class LiteLLM_OrganizationTableUpdate(LiteLLMPydanticObjectBase):
+class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable):
     """Represents user-controllable params for a LiteLLM_OrganizationTable record"""
 
     organization_id: Optional[str] = None
diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py
index 309438a17ff..c3fca196d57 100644
--- a/litellm/proxy/management_endpoints/organization_endpoints.py
+++ b/litellm/proxy/management_endpoints/organization_endpoints.py
@@ -36,6 +36,28 @@ from litellm.proxy.utils import PrismaClient
 router = APIRouter()
 
 
+def handle_nested_budget_structure_in_organization_update_request(raw_data: dict) -> dict:
+    """
+    Transform organization update request to handle UI payload format.
+    
+    The UI sends nested budget data in 'litellm_budget_table', but our
+    model expects flat budget fields at the top level.
+    """
+    transformed_data = raw_data.copy()
+    
+    # Handle nested budget structure from UI
+    if 'litellm_budget_table' in transformed_data:
+        budget_data = transformed_data.pop('litellm_budget_table', {})
+        if budget_data:
+            # Extract valid budget fields and merge into top level
+            budget_fields = LiteLLM_BudgetTable.model_fields.keys()
+            for key, value in budget_data.items():
+                if key in budget_fields and value is not None:
+                    transformed_data[key] = value
+    
+    return transformed_data
+
+
 @router.post(
     "/organization/new",
     tags=["organization management"],
@@ -248,7 +270,7 @@ async def _set_object_permission(
     response_model=LiteLLM_OrganizationTableWithMembers,
 )
 async def update_organization(
-    data: LiteLLM_OrganizationTableUpdate,
+    request: Request,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
@@ -270,6 +292,13 @@ async def update_organization(
             },
         )
 
+    # Transform UI payload to expected format
+    raw_data = await request.json()
+    raw_data_with_flat_budget_fields = handle_nested_budget_structure_in_organization_update_request(raw_data)
+    
+    # Create validated data model
+    data = LiteLLM_OrganizationTableUpdate(**raw_data_with_flat_budget_fields)
+    
     if data.updated_by is None:
         data.updated_by = user_api_key_dict.user_id
 
@@ -293,6 +322,23 @@ async def update_organization(
             existing_organization_row=existing_organization_row,
         )
 
+    # Handle budget updates if budget fields are provided
+    budget_fields = {k: v for k, v in data.model_dump().items() 
+                    if k in LiteLLM_BudgetTable.model_fields.keys() and v is not None}
+    
+    if budget_fields and existing_organization_row.budget_id:
+        await update_budget(
+            budget_obj=BudgetNewRequest(
+                budget_id=existing_organization_row.budget_id,
+                **budget_fields
+            ),
+            user_api_key_dict=user_api_key_dict,
+        )
+    
+    # Remove budget fields from organization update data
+    for field in LiteLLM_BudgetTable.model_fields.keys():
+        updated_organization_row.pop(field, None)
+
     response = await prisma_client.db.litellm_organizationtable.update(
         where={"organization_id": data.organization_id},
         data=updated_organization_row,
diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py
index c05e528846d..c176ef20374 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py
@@ -235,3 +235,38 @@ async def test_organization_update_object_permissions_missing_permission_record(
 
     # Verify upsert was called to create new record
     mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
+
+
+def test_transform_organization_update_request():
+    """
+    Test the request transformer handles nested litellm_budget_table structure from UI.
+    """
+    from litellm.proxy.management_endpoints.organization_endpoints import (
+        transform_organization_update_request,
+    )
+
+    # Test UI payload with nested budget structure
+    ui_payload = {
+        "organization_id": "test_org_id",
+        "organization_alias": "Test Org",
+        "litellm_budget_table": {
+            "max_budget": "1000.0",
+            "tpm_limit": 5000,
+            "rpm_limit": None
+        }
+    }
+    
+    # Transform the payload
+    transformed = transform_organization_update_request(ui_payload)
+    
+    # Verify nested structure was flattened to top-level fields
+    assert transformed["max_budget"] == "1000.0"
+    assert transformed["tpm_limit"] == 5000
+    assert "rpm_limit" not in transformed  # None values shouldn't be included
+    
+    # Verify nested field was removed
+    assert "litellm_budget_table" not in transformed
+    
+    # Verify other fields preserved
+    assert transformed["organization_id"] == "test_org_id"
+    assert transformed["organization_alias"] == "Test Org"