From 6cd19042c0d6f740421261768909cc991ce2f98e Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Sat, 11 Oct 2025 10:58:50 -0700
Subject: [PATCH] fix org endpoints

---
 .../organization_endpoints.py                 | 64 +++++++++++--------
 1 file changed, 37 insertions(+), 27 deletions(-)

diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py
index a8cbf6209af..c477c57a1fc 100644
--- a/litellm/proxy/management_endpoints/organization_endpoints.py
+++ b/litellm/proxy/management_endpoints/organization_endpoints.py
@@ -29,7 +29,6 @@ from litellm.proxy.management_helpers.object_permission_utils import (
 )
 from litellm.proxy.management_helpers.utils import (
     get_new_internal_user_defaults,
-    handle_budget_for_entity,
     management_endpoint_wrapper,
 )
 from litellm.proxy.utils import PrismaClient
@@ -169,17 +168,30 @@ async def new_organization(
         except Exception:
             pass
 
-    # Handle budget creation/assignment using common helper
-    budget_id = await handle_budget_for_entity(
-        data=data,
-        existing_budget_id=None,
-        user_api_key_dict=user_api_key_dict,
-        prisma_client=prisma_client,
-        litellm_proxy_admin_name=litellm_proxy_admin_name,
-    )
-    
-    if budget_id is not None:
-        data.budget_id = budget_id
+    if data.budget_id is None:
+        """
+        Every organization needs a budget attached.
+
+        If none provided, create one based on provided values
+        """
+        budget_params = LiteLLM_BudgetTable.model_fields.keys()
+
+        # Only include Budget Params when creating an entry in litellm_budgettable
+        _json_data = data.json(exclude_none=True)
+        _budget_data = {k: v for k, v in _json_data.items() if k in budget_params}
+        budget_row = LiteLLM_BudgetTable(**_budget_data)
+
+        new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
+
+        _budget = await prisma_client.db.litellm_budgettable.create(
+            data={
+                **new_budget,  # type: ignore
+                "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+                "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+            }
+        )  # type: ignore
+
+        data.budget_id = _budget.budget_id
 
     ## Handle Object Permission - MCP, Vector Stores etc.
     object_permission_id = await _set_object_permission(
@@ -310,22 +322,20 @@ async def update_organization(
             existing_organization_row=existing_organization_row,
         )
 
-    from litellm.proxy.proxy_server import litellm_proxy_admin_name
-
-    # Handle budget updates using common helper
-    budget_id = await handle_budget_for_entity(
-        data=data,
-        existing_budget_id=existing_organization_row.budget_id,
-        user_api_key_dict=user_api_key_dict,
-        prisma_client=prisma_client,
-        litellm_proxy_admin_name=litellm_proxy_admin_name,
-    )
+    # 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}
     
-    # Update budget_id if it changed
-    if budget_id != existing_organization_row.budget_id:
-        updated_organization_row["budget_id"] = budget_id
+    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 (they're handled via budget table)
+    # Remove budget fields from organization update data
     for field in LiteLLM_BudgetTable.model_fields.keys():
         updated_organization_row.pop(field, None)
 
@@ -957,4 +967,4 @@ async def add_member_to_organization(
     except Exception as e:
         raise ValueError(
             f"Error adding member={member} to organization={organization_id}: {e}"
-        )
+        )
\ No newline at end of file