fix(proxy_server.py): allow mapping a user to an org

litellm/proxy/_types.py

@@ -363,6 +363,8 @@ class NewUserRequest(GenerateKeyRequest):
     max_budget: Optional[float] = None
     user_email: Optional[str] = None
     user_role: Optional[str] = None
+    teams: Optional[list] = None
+    organization_id: Optional[str] = None
     auto_create_key: bool = (
         True  # flag used for returning a key as part of the /user/new response
     )

litellm/proxy/proxy_server.py

@@ -2553,6 +2553,8 @@ async def generate_key_helper_fn(
     allowed_cache_controls: Optional[list] = [],
     permissions: Optional[dict] = {},
     model_max_budget: Optional[dict] = {},
+    teams: Optional[list] = None,
+    organization_id: Optional[str] = None,
     table_name: Optional[Literal["key", "user"]] = None,
 ):
     global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name
@@ -2600,6 +2602,7 @@ async def generate_key_helper_fn(
             "user_email": user_email,
             "user_id": user_id,
             "team_id": team_id,
+            "organization_id": organization_id,
             "user_role": user_role,
             "spend": spend,
             "models": models,
@@ -2610,6 +2613,8 @@ async def generate_key_helper_fn(
             "budget_reset_at": reset_at,
             "allowed_cache_controls": allowed_cache_controls,
         }
+        if teams is not None:
+            user_data["teams"] = teams
         key_data = {
             "token": token,
             "key_alias": key_alias,
@@ -2710,6 +2715,8 @@ async def generate_key_helper_fn(
             await custom_db_client.insert_data(value=key_data, table_name="key")
     except Exception as e:
         traceback.print_exc()
+        if isinstance(e, HTTPException):
+            raise e
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail={"error": "Internal Server Error."},
@@ -5470,6 +5477,8 @@ async def new_user(data: NewUserRequest):
     Parameters:
     - user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated.
     - user_alias: Optional[str] - A descriptive name for you to know who this user id refers to.
+    - teams: Optional[list] - specify a list of team id's a user belongs to.
+    - organization_id: Optional[str] - specify the org a user belongs to.
     - user_email: Optional[str] - Specify a user email.
     - user_role: Optional[str] - Specify a user role - "admin", "app_owner", "app_user"
     - max_budget: Optional[float] - Specify max budget for a given user.

litellm/proxy/schema.prisma

@@ -53,6 +53,7 @@ model LiteLLM_OrganizationTable {
     updated_by String
     litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
     teams LiteLLM_TeamTable[] 
+    users LiteLLM_UserTable[]
 }
 
 // Model info for teams, just has model aliases for now.
@@ -99,6 +100,7 @@ model LiteLLM_UserTable {
 		user_id    String @id
     user_alias String? 
     team_id    String?
+    organization_id String?
     teams    String[] @default([])
     user_role  String?
 		max_budget Float?
@@ -113,6 +115,7 @@ model LiteLLM_UserTable {
     allowed_cache_controls String[] @default([])
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
+    litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
 }
 
 // Generate Tokens for Proxy

litellm/proxy/utils.py

@@ -461,7 +461,12 @@ class ProxyLogging:
         """
         ### ALERTING ###
         if isinstance(original_exception, HTTPException):
-            error_message = original_exception.detail
+            if isinstance(original_exception.detail, str):
+                error_message = original_exception.detail
+            elif isinstance(original_exception.detail, dict):
+                error_message = json.dumps(original_exception.detail)
+            else:
+                error_message = str(original_exception)
         else:
             error_message = str(original_exception)
         if isinstance(traceback_str, str):
@@ -1159,13 +1164,26 @@ class PrismaClient:
                 return new_verification_token
             elif table_name == "user":
                 db_data = self.jsonify_object(data=data)
-                new_user_row = await self.db.litellm_usertable.upsert(
+                try:
-                    where={"user_id": data["user_id"]},
+                    new_user_row = await self.db.litellm_usertable.upsert(
-                    data={
+                        where={"user_id": data["user_id"]},
-                        "create": {**db_data},  # type: ignore
+                        data={
-                        "update": {},  # don't do anything if it already exists
+                            "create": {**db_data},  # type: ignore
-                    },
+                            "update": {},  # don't do anything if it already exists
-                )
+                        },
+                    )
+                except Exception as e:
+                    if (
+                        "Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"
+                        in str(e)
+                    ):
+                        raise HTTPException(
+                            status_code=400,
+                            detail={
+                                "error": f"Foreign Key Constraint failed. Organization ID={db_data['organization_id']} does not exist in LiteLLM_OrganizationTable"
+                            },
+                        )
+                    raise e
                 verbose_proxy_logger.info("Data Inserted into User Table")
                 return new_user_row
             elif table_name == "team":

schema.prisma

@@ -53,6 +53,7 @@ model LiteLLM_OrganizationTable {
     updated_by String
     litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
     teams LiteLLM_TeamTable[] 
+    users LiteLLM_UserTable[]
 }
 
 // Model info for teams, just has model aliases for now.
@@ -99,6 +100,7 @@ model LiteLLM_UserTable {
 		user_id    String @id
     user_alias String? 
     team_id    String?
+    organization_id String?
     teams    String[] @default([])
     user_role  String?
 		max_budget Float?
@@ -113,6 +115,7 @@ model LiteLLM_UserTable {
     allowed_cache_controls String[] @default([])
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
+    litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
 }
 
 // Generate Tokens for Proxy