diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5c9cc6c19..26a0438c6 100644 --- a/litellm/proxy/_types.py +++ b/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 ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 74d84a456..6b2a5d21e 100644 --- a/litellm/proxy/proxy_server.py +++ b/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. diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 0a490451e..5ec73c9dc 100644 --- a/litellm/proxy/schema.prisma +++ b/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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a36124792..fd302c52e 100644 --- a/litellm/proxy/utils.py +++ b/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( - where={"user_id": data["user_id"]}, - data={ - "create": {**db_data}, # type: ignore - "update": {}, # don't do anything if it already exists - }, - ) + try: + new_user_row = await self.db.litellm_usertable.upsert( + where={"user_id": data["user_id"]}, + data={ + "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": diff --git a/schema.prisma b/schema.prisma index 0a490451e..5ec73c9dc 100644 --- a/schema.prisma +++ b/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