phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(feat proxy) [beta] add support for organization role based access controls (#6112)

Browse source 
* track LiteLLM_OrganizationMembership

* add add_internal_user_to_organization

* add org membership to schema

* read organization membership when reading user info in auth checks

* add check for valid organization_id

* add test for test_create_new_user_in_organization

* test test_create_new_user_in_organization

* add new ADMIN role

* add test for org admins creating teams

* add test for test_org_admin_create_user_permissions

* test_org_admin_create_user_team_wrong_org_permissions

* test_org_admin_create_user_team_wrong_org_permissions

* fix organization_role_based_access_check

* fix getting user members

* fix TeamBase

* fix types used for use role

* fix type checks

* sync prisma schema

* docs - organization admins

* fix use organization_endpoints for /organization management

* add types for org member endpoints

* fix role name for org admin

* add type for member add response

* add organization/member_add

* add error handling for adding members to an org

* add nice doc string for oranization/member_add

* fix test_create_new_user_in_organization

* linting fix

* use simple route changes

* fix types

* add organization member roles

* add org admin auth checks

* add auth checks for orgs

* test for creating teams as org admin

* simplify org id usage

* fix typo

* test test_org_admin_create_user_team_wrong_org_permissions

* fix type check issue

* code quality fix

* fix schema.prisma
This commit is contained in:
Ishaan Jaff 2024-10-09 15:18:18 +05:30 • committed by GitHub
parent 945267a511
commit 1fd437e263
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1474 additions and 261 deletions
Download patch file Download diff file Expand all files Collapse all files

145
docs/my-website/docs/proxy/access_control.md Normal file
View file

@ -0,0 +1,145 @@
# Role-based Access Controls (RBAC)
Role-based access control (RBAC) is based on Organizations, Teams and Internal User Roles
- `Organizations` are the top-level entities that contain Teams.
- `Team` - A Team is a collection of multiple `Internal Users`
- `Internal Users` - users that can create keys, make LLM API calls, view usage on LiteLLM
- `Roles` define the permissions of an `Internal User`
- `Virtual Keys` - Keys are used for authentication to the LiteLLM API. Keys are tied to a `Internal User` and `Team`
## Roles
**Admin Roles**
- `proxy_admin`: admin over the platform
- `proxy_admin_viewer`: can login, view all keys, view all spend. **Cannot** create/delete keys, add new users.
**Organization Roles**
- `organization_admin`: admin over the organization. Can create teams and users within their organization
**Internal User Roles**
- `internal_user`: can login, view/create/delete their own keys, view their spend. **Cannot** add new users.
- `internal_user_viewer`: can login, view their own keys, view their own spend. **Cannot** create/delete keys, add new users.
## Managing Organizations
### 1. Creating a new Organization
Any user with role=`proxy_admin` can create a new organization
**Usage**
[**API Reference for /organization/new**](https://litellm-api.up.railway.app/#/organization%20management/new_organization_organization_new_post)
```shell
curl --location 'http://0.0.0.0:4000/organization/new' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"organization_alias": "marketing_department",
"models": ["gpt-4"],
"max_budget": 20
}'
```
Expected Response
```json
{
"organization_id": "ad15e8ca-12ae-46f4-8659-d02debef1b23",
"organization_alias": "marketing_department",
"budget_id": "98754244-3a9c-4b31-b2e9-c63edc8fd7eb",
"metadata": {},
"models": [
"gpt-4"
],
"created_by": "109010464461339474872",
"updated_by": "109010464461339474872",
"created_at": "2024-10-08T18:30:24.637000Z",
"updated_at": "2024-10-08T18:30:24.637000Z"
}
```
### 2. Adding an `organization_admin` to an Organization
Create a user (ishaan@berri.ai) as an `organization_admin` for the `marketing_department` Organization (from [step 1](#1-creating-a-new-organization))
Users with the following roles can call `/organization/member_add`
- `proxy_admin`
- `organization_admin` only within their own organization
```shell
curl -X POST 'http://0.0.0.0:4000/organization/member_add' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{"organization_id": "ad15e8ca-12ae-46f4-8659-d02debef1b23", "member": {"role": "organization_admin", "user_id": "ishaan@berri.ai"}}'
```
Now a user with user_id = `ishaan@berri.ai` and role = `organization_admin` has been created in the `marketing_department` Organization
Create a Virtual Key for user_id = `ishaan@berri.ai`. The User can then use the Virtual key for their Organization Admin Operations
```shell
curl --location 'http://0.0.0.0:4000/key/generate' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"user_id": "ishaan@berri.ai"
}'
```
Expected Response
```json
{
"models": [],
"user_id": "ishaan@berri.ai",
"key": "sk-7shH8TGMAofR4zQpAAo6kQ",
"key_name": "sk-...o6kQ",
}
```
### 3. `Organization Admin` - Create a Team
The organization admin will use the virtual key created in [step 2](#2-adding-an-organization_admin-to-an-organization) to create a `Team` within the `marketing_department` Organization
```shell
curl --location 'http://0.0.0.0:4000/team/new' \
--header 'Authorization: Bearer sk-7shH8TGMAofR4zQpAAo6kQ' \
--header 'Content-Type: application/json' \
--data '{
"team_alias": "engineering_team",
"organization_id": "ad15e8ca-12ae-46f4-8659-d02debef1b23",
}'
```
This will create the team `engineering_team` within the `marketing_department` Organization
Expected Response
```json
{
"team_alias": "engineering_team",
"team_id": "01044ee8-441b-45f4-be7d-c70e002722d8",
"organization_id": "ad15e8ca-12ae-46f4-8659-d02debef1b23",
}
```
### `Organization Admin` - Add an `Internal User`
The organization admin will use the virtual key created in [step 2](#2-adding-an-organization_admin-to-an-organization) to add an Internal User to the `engineering_team` Team.
- We will assign role=`internal_user` so the user can create Virtual Keys for themselves
- `team_id` is from [step 3](#3-organization-admin---create-a-team)
```shell
curl -X POST 'http://0.0.0.0:4000/team/member_add' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{"team_id": "01044ee8-441b-45f4-be7d-c70e002722d8",, "member": {"role": "internal_user", "user_id": "krrish@berri.ai"}}'
```

7
docs/my-website/sidebars.js
View file

@ -78,7 +78,12 @@ const sidebars = {
{
type: "category",
label: "Admin UI",
items: ["proxy/ui", "proxy/self_serve", "proxy/custom_sso"],
items: [
"proxy/ui",
"proxy/self_serve",
"proxy/access_control",
"proxy/custom_sso"
],
},
{
type: "category",

189
litellm/proxy/_types.py
View file

@ -34,6 +34,7 @@ class LitellmUserRoles(str, enum.Enum):
Admin Roles:
PROXY_ADMIN: admin over the platform
PROXY_ADMIN_VIEW_ONLY: can login, view all own keys, view all spend
ORG_ADMIN: admin over a specific organization, can create teams, users only within their organization
Internal User Roles:
INTERNAL_USER: can login, view/create/delete their own keys, view their spend
@ -53,6 +54,9 @@ class LitellmUserRoles(str, enum.Enum):
PROXY_ADMIN = "proxy_admin"
PROXY_ADMIN_VIEW_ONLY = "proxy_admin_viewer"
# Organization admins
ORG_ADMIN = "org_admin"
# Internal User Roles
INTERNAL_USER = "internal_user"
INTERNAL_USER_VIEW_ONLY = "internal_user_viewer"
@ -359,6 +363,20 @@ class LiteLLMRoutes(enum.Enum):
"/team/member_delete",
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##
# Routes only an Org Admin Can Access
org_admin_only_routes = [
"/organization/info",
"/organization/delete",
"/organization/member_add",
]
# All routes accesible by an Org Admin
org_admin_allowed_routes = (
org_admin_only_routes + management_routes + self_managed_routes
)
# class LiteLLMAllowedRoutes(LiteLLMBase):
# """
@ -695,12 +713,9 @@ class NewUserRequest(_GenerateKeyRequest):
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.TEAM,
LitellmUserRoles.CUSTOMER,
]
] = 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
)
@ -716,12 +731,9 @@ class NewUserResponse(GenerateKeyResponse):
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.TEAM,
LitellmUserRoles.CUSTOMER,
]
] = None
teams: Optional[list] = None
organization_id: Optional[str] = None
user_alias: Optional[str] = None
@ -739,8 +751,6 @@ class UpdateUserRequest(GenerateRequestBase):
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.TEAM,
LitellmUserRoles.CUSTOMER,
]
] = None
max_budget: Optional[float] = None
@ -811,7 +821,14 @@ class DeleteCustomerRequest(LiteLLMBase):
class Member(LiteLLMBase):
role: Literal["admin", "user"]
role: Literal[
LitellmUserRoles.ORG_ADMIN,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
# older Member roles
"admin",
"user",
]
user_id: Optional[str] = None
user_email: Optional[str] = None
@ -857,51 +874,6 @@ class GlobalEndUsersSpend(LiteLLMBase):
endTime: Optional[datetime] = None
class TeamMemberAddRequest(LiteLLMBase):
team_id: str
member: Union[List[Member], Member]
max_budget_in_team: Optional[float] = None # Users max budget within the team
def __init__(self, **data):
member_data = data.get("member")
if isinstance(member_data, list):
# If member is a list of dictionaries, convert each dictionary to a Member object
members = [Member(**item) for item in member_data]
# Replace member_data with the list of Member objects
data["member"] = members
elif isinstance(member_data, dict):
# If member is a dictionary, convert it to a single Member object
member = Member(**member_data)
# Replace member_data with the single Member object
data["member"] = member
# Call the superclass __init__ method to initialize the object
super().__init__(**data)
class TeamMemberDeleteRequest(LiteLLMBase):
team_id: str
user_id: Optional[str] = None
user_email: Optional[str] = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values):
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
max_budget_in_team: float
class TeamMemberUpdateResponse(LiteLLMBase):
team_id: str
user_id: str
user_email: Optional[str] = None
max_budget_in_team: float
class UpdateTeamRequest(LiteLLMBase):
"""
UpdateTeamRequest, used by /team/update when you need to update a team
@ -1444,6 +1416,26 @@ class LiteLLM_Config(LiteLLMBase):
param_value: Dict
class LiteLLM_OrganizationMembershipTable(LiteLLMBase):
"""
This is the table that track what organizations a user belongs to and users spend within the organization
"""
user_id: str
organization_id: str
user_role: Optional[str] = None
spend: float = 0.0
budget_id: Optional[str] = None
created_at: datetime
updated_at: datetime
user: Optional[Any] = (
None # You might want to replace 'Any' with a more specific type if available
)
litellm_budget_table: Optional[LiteLLM_BudgetTable] = None
model_config = ConfigDict(protected_namespaces=())
class LiteLLM_UserTable(LiteLLMBase):
user_id: str
max_budget: Optional[float]
@ -1455,6 +1447,7 @@ class LiteLLM_UserTable(LiteLLMBase):
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
user_role: Optional[str] = None
organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None
@model_validator(mode="before")
@classmethod
@ -1907,11 +1900,99 @@ class LiteLLM_TeamMembership(LiteLLMBase):
litellm_budget_table: Optional[LiteLLM_BudgetTable]
#### Organization / Team Member Requests ####
class MemberAddRequest(LiteLLMBase):
member: Union[List[Member], Member]
def __init__(self, **data):
member_data = data.get("member")
if isinstance(member_data, list):
# If member is a list of dictionaries, convert each dictionary to a Member object
members = [Member(**item) for item in member_data]
# Replace member_data with the list of Member objects
data["member"] = members
elif isinstance(member_data, dict):
# If member is a dictionary, convert it to a single Member object
member = Member(**member_data)
# Replace member_data with the single Member object
data["member"] = member
# Call the superclass __init__ method to initialize the object
super().__init__(**data)
class TeamAddMemberResponse(LiteLLM_TeamTable):
updated_users: List[LiteLLM_UserTable]
updated_team_memberships: List[LiteLLM_TeamMembership]
class OrganizationAddMemberResponse(LiteLLMBase):
organization_id: str
updated_users: List[LiteLLM_UserTable]
updated_organization_memberships: List[LiteLLM_OrganizationMembershipTable]
class MemberDeleteRequest(LiteLLMBase):
user_id: Optional[str] = None
user_email: Optional[str] = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values):
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
class MemberUpdateResponse(LiteLLMBase):
user_id: str
user_email: Optional[str] = None
# Team Member Requests
class TeamMemberAddRequest(MemberAddRequest):
team_id: str
max_budget_in_team: Optional[float] = None # Users max budget within the team
class TeamMemberDeleteRequest(MemberDeleteRequest):
team_id: str
class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
max_budget_in_team: float
class TeamMemberUpdateResponse(MemberUpdateResponse):
team_id: str
max_budget_in_team: float
# Organization Member Requests
class OrganizationMemberAddRequest(MemberAddRequest):
organization_id: str
max_budget_in_organization: Optional[float] = (
None # Users max budget within the organization
)
class OrganizationMemberDeleteRequest(MemberDeleteRequest):
organization_id: str
class OrganizationMemberUpdateRequest(OrganizationMemberDeleteRequest):
max_budget_in_organization: float
class OrganizationMemberUpdateResponse(MemberUpdateResponse):
organization_id: str
max_budget_in_organization: float
##########################################
class TeamInfoResponseObject(TypedDict):
team_id: str
team_info: LiteLLM_TeamTable

31
litellm/proxy/auth/auth_checks.py
View file

@ -32,6 +32,8 @@ from litellm.proxy.auth.route_checks import is_llm_api_route
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_to_opentelemetry
from litellm.types.services import ServiceLoggerPayload, ServiceTypes
from .auth_checks_organization import organization_role_based_access_check
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -63,6 +65,7 @@ def common_checks(
7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
8. [OPTIONAL] If guardrails modified - is request allowed to change this
9. Check if request body is safe
10. [OPTIONAL] Organization checks - is user_object.organization_id is set, run these checks
"""
_model = request_body.get("model", None)
if team_object is not None and team_object.blocked is True:
@ -73,6 +76,7 @@ def common_checks(
if (
_model is not None
and team_object is not None
and team_object.models is not None
and len(team_object.models) > 0
and _model not in team_object.models
):
@ -202,6 +206,12 @@ def common_checks(
"error": "Your team does not have permission to modify guardrails."
},
)
# 10 [OPTIONAL] Organization RBAC checks
organization_role_based_access_check(
user_object=user_object, route=route, request_body=request_body
)
return True
@ -403,17 +413,30 @@ async def get_user_object(
try:
response = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
where={"user_id": user_id}, include={"organization_memberships": True}
)
if response is None:
if user_id_upsert:
response = await prisma_client.db.litellm_usertable.create(
data={"user_id": user_id}
data={"user_id": user_id},
include={"organization_memberships": True},
)
else:
raise Exception
if (
response.organization_memberships is not None
and len(response.organization_memberships) > 0
):
# dump each organization membership to type LiteLLM_OrganizationMembershipTable
_dumped_memberships = [
membership.model_dump()
for membership in response.organization_memberships
if membership is not None
]
response.organization_memberships = _dumped_memberships
_response = LiteLLM_UserTable(**dict(response))
response_dict = _response.model_dump()
@ -421,9 +444,9 @@ async def get_user_object(
await user_api_key_cache.async_set_cache(key=user_id, value=response_dict)
return _response
except Exception: # if user not in db
except Exception as e: # if user not in db
raise ValueError(
f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call."
f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {e}"
)

161
litellm/proxy/auth/auth_checks_organization.py Normal file
View file

@ -0,0 +1,161 @@
"""
Auth Checks for Organizations
"""
from typing import Dict, List, Optional, Tuple
from fastapi import status
from litellm.proxy._types import *
def organization_role_based_access_check(
request_body: dict,
user_object: Optional[LiteLLM_UserTable],
route: str,
):
"""
Role based access control checks only run if a user is part of an Organization
Organization Checks:
ONLY RUN IF user_object.organization_memberships is not None
1. Only Proxy Admins can access /organization/new
2. IF route is a LiteLLMRoutes.org_admin_only_routes, then check if user is an Org Admin for that organization
"""
if user_object is None:
return
passed_organization_id: Optional[str] = request_body.get("organization_id", None)
if route == "/organization/new":
if user_object.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise ProxyException(
message=f"Only proxy admins can create new organizations. You are {user_object.user_role}",
type=ProxyErrorTypes.auth_error.value,
param="user_role",
code=status.HTTP_401_UNAUTHORIZED,
)
if user_object.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
# Checks if route is an Org Admin Only Route
if route in LiteLLMRoutes.org_admin_only_routes.value:
_user_organizations, _user_organization_role_mapping = (
get_user_organization_info(user_object)
)
if user_object.organization_memberships is None:
raise ProxyException(
message=f"Tried to access route={route} but you are not a member of any organization. Please contact the proxy admin to request access.",
type=ProxyErrorTypes.auth_error.value,
param="organization_id",
code=status.HTTP_401_UNAUTHORIZED,
)
if passed_organization_id is None:
raise ProxyException(
message="Passed organization_id is None, please pass an organization_id in your request",
type=ProxyErrorTypes.auth_error.value,
param="organization_id",
code=status.HTTP_401_UNAUTHORIZED,
)
user_role: Optional[LitellmUserRoles] = _user_organization_role_mapping.get(
passed_organization_id
)
if user_role is None:
raise ProxyException(
message=f"You do not have a role within the selected organization. Passed organization_id: {passed_organization_id}. Please contact the organization admin to request access.",
type=ProxyErrorTypes.auth_error.value,
param="organization_id",
code=status.HTTP_401_UNAUTHORIZED,
)
if user_role != LitellmUserRoles.ORG_ADMIN.value:
raise ProxyException(
message=f"You do not have the required role to perform {route} in Organization {passed_organization_id}. Your role is {user_role} in Organization {passed_organization_id}",
type=ProxyErrorTypes.auth_error.value,
param="user_role",
code=status.HTTP_401_UNAUTHORIZED,
)
elif route == "/team/new":
# if user is part of multiple teams, then they need to specify the organization_id
_user_organizations, _user_organization_role_mapping = (
get_user_organization_info(user_object)
)
if (
user_object.organization_memberships is not None
and len(user_object.organization_memberships) > 0
):
if passed_organization_id is None:
raise ProxyException(
message=f"Passed organization_id is None, please specify the organization_id in your request. You are part of multiple organizations: {_user_organizations}",
type=ProxyErrorTypes.auth_error.value,
param="organization_id",
code=status.HTTP_401_UNAUTHORIZED,
)
_user_role_in_passed_org = _user_organization_role_mapping.get(
passed_organization_id
)
if _user_role_in_passed_org != LitellmUserRoles.ORG_ADMIN.value:
raise ProxyException(
message=f"You do not have the required role to call {route}. Your role is {_user_role_in_passed_org} in Organization {passed_organization_id}",
type=ProxyErrorTypes.auth_error.value,
param="user_role",
code=status.HTTP_401_UNAUTHORIZED,
)
def get_user_organization_info(
user_object: LiteLLM_UserTable,
) -> Tuple[List[str], Dict[str, Optional[LitellmUserRoles]]]:
"""
Helper function to extract user organization information.
Args:
user_object (LiteLLM_UserTable): The user object containing organization memberships.
Returns:
Tuple[List[str], Dict[str, Optional[LitellmUserRoles]]]: A tuple containing:
- List of organization IDs the user is a member of
- Dictionary mapping organization IDs to user roles
"""
_user_organizations: List[str] = []
_user_organization_role_mapping: Dict[str, Optional[LitellmUserRoles]] = {}
if user_object.organization_memberships is not None:
for _membership in user_object.organization_memberships:
if _membership.organization_id is not None:
_user_organizations.append(_membership.organization_id)
_user_organization_role_mapping[_membership.organization_id] = _membership.user_role # type: ignore
return _user_organizations, _user_organization_role_mapping
def _user_is_org_admin(
request_data: dict,
user_object: Optional[LiteLLM_UserTable] = None,
) -> bool:
"""
Helper function to check if user is an org admin for the passed organization_id
"""
if request_data.get("organization_id", None) is None:
return False
if user_object is None:
return False
if user_object.organization_memberships is None:
return False
for _membership in user_object.organization_memberships:
if _membership.organization_id == request_data.get("organization_id", None):
if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value:
return True
return False

10
litellm/proxy/auth/route_checks.py
View file

@ -13,10 +13,11 @@ from litellm.proxy._types import (
)
from litellm.proxy.utils import hash_token
from .auth_checks_organization import _user_is_org_admin
from .auth_utils import _has_user_setup_sso
def non_admin_allowed_routes_check(
def non_proxy_admin_allowed_routes_check(
user_obj: Optional[LiteLLM_UserTable],
_user_role: Optional[LitellmUserRoles],
route: str,
@ -26,7 +27,7 @@ def non_admin_allowed_routes_check(
request_data: dict,
):
"""
Checks if Non-Admin User is allowed to access the route
Checks if Non Proxy Admin User is allowed to access the route
"""
# Check user has defined custom admin routes
@ -106,6 +107,11 @@ def non_admin_allowed_routes_check(
and route in LiteLLMRoutes.internal_user_routes.value
):
pass
elif (
_user_is_org_admin(request_data=request_data, user_object=user_obj)
and route in LiteLLMRoutes.org_admin_allowed_routes.value
):
pass
elif (
_user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
and route in LiteLLMRoutes.internal_user_view_only_routes.value

4
litellm/proxy/auth/user_api_key_auth.py
View file

@ -69,7 +69,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.oauth2_check import check_oauth2_token
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.route_checks import non_admin_allowed_routes_check
from litellm.proxy.auth.route_checks import non_proxy_admin_allowed_routes_check
from litellm.proxy.auth.service_account_checks import service_account_checks
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.utils import _to_ns
@ -1042,7 +1042,7 @@ async def user_api_key_auth(
_user_role = _get_user_role(user_obj=user_obj)
if not _is_user_proxy_admin(user_obj=user_obj): # if non-admin
non_admin_allowed_routes_check(
non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=_user_role,
route=route,

60
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -63,7 +63,6 @@ async def new_user(
- 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.
- send_invite_email: Optional[bool] - Specify if an invite email should be sent.
- user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20`
@ -79,6 +78,18 @@ async def new_user(
- expires: (datetime) Datetime object for when key expires.
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
- max_budget: (float|None) Max budget for given user.
Usage Example
```shell
curl -X POST "http://localhost:4000/user/new" \
-H "Content-Type: application/json" \
-H "Authorization: Bearer sk-1234" \
-d '{
"username": "new_user",
"email": "new_user@example.com"
}'
```
"""
from litellm.proxy.proxy_server import general_settings, proxy_logging_obj
@ -106,6 +117,7 @@ async def new_user(
response = await generate_key_helper_fn(request_type="user", **data_json)
# Admin UI Logic
# Add User to Team and Organization
# if team_id passed add this user to the team
if data_json.get("team_id", None) is not None:
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
@ -888,3 +900,49 @@ async def delete_user(
)
return deleted_users
async def add_internal_user_to_organization(
user_id: str,
organization_id: str,
user_role: LitellmUserRoles,
):
"""
Helper function to add an internal user to an organization
Adds the user to LiteLLM_OrganizationMembership table
- Checks if organization_id exists
Raises:
- Exception if database not connected
- Exception if user_id or organization_id not found
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise Exception("Database not connected")
try:
# Check if organization_id exists
organization_row = await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": organization_id}
)
if organization_row is None:
raise Exception(
f"Organization not found, passed organization_id={organization_id}"
)
# Create a new organization membership entry
new_membership = await prisma_client.db.litellm_organizationmembership.create(
data={
"user_id": user_id,
"organization_id": organization_id,
"user_role": user_role,
# Note: You can also set budget within an organization if needed
}
)
return new_membership
except Exception as e:
raise Exception(f"Failed to add user to organization: {str(e)}")

433
litellm/proxy/management_endpoints/organization_endpoints.py Normal file
View file

@ -0,0 +1,433 @@
"""
Endpoints for /organization operations
/organization/new
/organization/update
/organization/delete
/organization/info
"""
#### ORGANIZATION MANAGEMENT ####
import asyncio
import copy
import json
import re
import secrets
import traceback
import uuid
from datetime import datetime, timedelta, timezone
from typing import List, Optional, Tuple
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_helpers.utils import (
get_new_internal_user_defaults,
management_endpoint_wrapper,
)
from litellm.proxy.utils import PrismaClient
from litellm.secret_managers.main import get_secret
router = APIRouter()
@router.post(
"/organization/new",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
response_model=NewOrganizationResponse,
)
async def new_organization(
data: NewOrganizationRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Allow orgs to own teams
Set org level budgets + model access.
Only admins can create orgs.
# Parameters
- `organization_alias`: *str* = The name of the organization.
- `models`: *List* = The models the organization has access to.
- `budget_id`: *Optional[str]* = The id for a budget (tpm/rpm/max budget) for the organization.
### IF NO BUDGET ID - CREATE ONE WITH THESE PARAMS ###
- `max_budget`: *Optional[float]* = Max budget for org
- `tpm_limit`: *Optional[int]* = Max tpm limit for org
- `rpm_limit`: *Optional[int]* = Max rpm limit for org
- `model_max_budget`: *Optional[dict]* = Max budget for a specific model
- `budget_duration`: *Optional[str]* = Frequency of reseting org budget
Case 1: Create new org **without** a budget_id
```bash
curl --location 'http://0.0.0.0:4000/organization/new' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"organization_alias": "my-secret-org",
"models": ["model1", "model2"],
"max_budget": 100
}'
```
Case 2: Create new org **with** a budget_id
```bash
curl --location 'http://0.0.0.0:4000/organization/new' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"organization_alias": "my-secret-org",
"models": ["model1", "model2"],
"budget_id": "428eeaa8-f3ac-4e85-a8fb-7dc8d7aa8689"
}'
```
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
if (
user_api_key_dict.user_role is None
or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
):
raise HTTPException(
status_code=401,
detail={
"error": f"Only admins can create orgs. Your role is = {user_api_key_dict.user_role}"
},
)
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
"""
Ensure only models that user has access to, are given to org
"""
if len(user_api_key_dict.models) == 0: # user has access to all models
pass
else:
if len(data.models) == 0:
raise HTTPException(
status_code=400,
detail={
"error": "User not allowed to give access to all models. Select models you want org to have access to."
},
)
for m in data.models:
if m not in user_api_key_dict.models:
raise HTTPException(
status_code=400,
detail={
"error": f"User not allowed to give access to model={m}. Models you have access to = {user_api_key_dict.models}"
},
)
organization_row = LiteLLM_OrganizationTable(
**data.json(exclude_none=True),
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,
)
new_organization_row = prisma_client.jsonify_object(
organization_row.json(exclude_none=True)
)
response = await prisma_client.db.litellm_organizationtable.create(
data={
**new_organization_row, # type: ignore
}
)
return response
@router.post(
"/organization/update",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_organization():
"""[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues"""
pass
@router.post(
"/organization/delete",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_organization():
"""[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues"""
pass
@router.post(
"/organization/info",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
)
async def info_organization(data: OrganizationRequest):
"""
Get the org specific information
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
if len(data.organizations) == 0:
raise HTTPException(
status_code=400,
detail={
"error": f"Specify list of organization id's to query. Passed in={data.organizations}"
},
)
response = await prisma_client.db.litellm_organizationtable.find_many(
where={"organization_id": {"in": data.organizations}},
include={"litellm_budget_table": True},
)
return response
@router.post(
"/organization/member_add",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
response_model=OrganizationAddMemberResponse,
)
@management_endpoint_wrapper
async def organization_member_add(
data: OrganizationMemberAddRequest,
http_request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> OrganizationAddMemberResponse:
"""
[BETA]
Add new members (either via user_email or user_id) to an organization
If user doesn't exist, new user row will also be added to User Table
Only proxy_admin or org_admin of organization, allowed to access this endpoint.
# Parameters:
- organization_id: str (required)
- member: Union[List[Member], Member] (required)
- role: Literal[LitellmUserRoles] (required)
- user_id: Optional[str]
- user_email: Optional[str]
Note: Either user_id or user_email must be provided for each member.
Example:
```
curl -X POST 'http://0.0.0.0:4000/organization/member_add' \
-H 'Authorization: Bearer sk-1234' \
-H 'Content-Type: application/json' \
-d '{
"organization_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849",
"member": {
"role": "internal_user",
"user_id": "krrish247652@berri.ai"
},
"max_budget_in_organization": 100.0
}'
```
The following is executed in this function:
1. Check if organization exists
2. Creates a new Internal User if the user_id or user_email is not found in LiteLLM_UserTable
3. Add Internal User to the `LiteLLM_OrganizationMembership` table
"""
try:
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
# Check if organization exists
existing_organization_row = (
await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": data.organization_id}
)
)
if existing_organization_row is None:
raise HTTPException(
status_code=404,
detail={
"error": f"Organization not found for organization_id={getattr(data, 'organization_id', None)}"
},
)
members: List[Member]
if isinstance(data.member, List):
members = data.member
else:
members = [data.member]
updated_users: List[LiteLLM_UserTable] = []
updated_organization_memberships: List[LiteLLM_OrganizationMembershipTable] = []
for member in members:
updated_user, updated_organization_membership = (
await add_member_to_organization(
member=member,
organization_id=data.organization_id,
prisma_client=prisma_client,
)
)
updated_users.append(updated_user)
updated_organization_memberships.append(updated_organization_membership)
return OrganizationAddMemberResponse(
organization_id=data.organization_id,
updated_users=updated_users,
updated_organization_memberships=updated_organization_memberships,
)
except Exception as e:
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({str(e)})"),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Authentication Error, " + str(e),
type=ProxyErrorTypes.auth_error,
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
async def add_member_to_organization(
member: Member,
organization_id: str,
prisma_client: PrismaClient,
) -> Tuple[LiteLLM_UserTable, LiteLLM_OrganizationMembershipTable]:
"""
Add a member to an organization
- Checks if member.user_id or member.user_email is in LiteLLM_UserTable
- If not found, create a new user in LiteLLM_UserTable
- Add user to organization in LiteLLM_OrganizationMembership
"""
try:
user_object: Optional[LiteLLM_UserTable] = None
existing_user_id_row = None
existing_user_email_row = None
## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable
if member.user_id is not None:
existing_user_id_row = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": member.user_id}
)
if member.user_email is not None:
existing_user_email_row = (
await prisma_client.db.litellm_usertable.find_unique(
where={"user_email": member.user_email}
)
)
## If user does not exist, create a new user
if existing_user_id_row is None and existing_user_email_row is None:
# Create a new user - since user does not exist
user_id: str = member.user_id or str(uuid.uuid4())
new_user_defaults = get_new_internal_user_defaults(
user_id=user_id,
user_email=member.user_email,
)
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore
if _returned_user is not None:
user_object = LiteLLM_UserTable(**_returned_user.model_dump())
elif existing_user_email_row is not None and len(existing_user_email_row) > 1:
raise HTTPException(
status_code=400,
detail={
"error": "Multiple users with this email found in db. Please use 'user_id' instead."
},
)
elif existing_user_email_row is not None:
user_object = LiteLLM_UserTable(**existing_user_email_row.model_dump())
elif existing_user_id_row is not None:
user_object = LiteLLM_UserTable(**existing_user_id_row.model_dump())
else:
raise HTTPException(
status_code=404,
detail={
"error": f"User not found for user_id={member.user_id} and user_email={member.user_email}"
},
)
if user_object is None:
raise ValueError(
f"User does not exist in LiteLLM_UserTable. user_id={member.user_id} and user_email={member.user_email}"
)
# Add user to organization
_organization_membership = (
await prisma_client.db.litellm_organizationmembership.create(
data={
"organization_id": organization_id,
"user_id": user_object.user_id,
"user_role": member.role,
}
)
)
organization_membership = LiteLLM_OrganizationMembershipTable(
**_organization_membership.model_dump()
)
return user_object, organization_membership
except Exception as e:
raise ValueError(f"Error adding member to organization: {e}")

206
litellm/proxy/proxy_server.py
View file

@ -185,6 +185,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
from litellm.proxy.management_endpoints.key_management_endpoints import (
router as key_management_router,
)
from litellm.proxy.management_endpoints.organization_endpoints import (
router as organization_router,
)
from litellm.proxy.management_endpoints.team_callback_endpoints import (
router as team_callback_router,
)
@ -6313,200 +6316,6 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
return
#### ORGANIZATION MANAGEMENT ####
@router.post(
"/organization/new",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
response_model=NewOrganizationResponse,
)
async def new_organization(
data: NewOrganizationRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Allow orgs to own teams
Set org level budgets + model access.
Only admins can create orgs.
# Parameters
- `organization_alias`: *str* = The name of the organization.
- `models`: *List* = The models the organization has access to.
- `budget_id`: *Optional[str]* = The id for a budget (tpm/rpm/max budget) for the organization.
### IF NO BUDGET ID - CREATE ONE WITH THESE PARAMS ###
- `max_budget`: *Optional[float]* = Max budget for org
- `tpm_limit`: *Optional[int]* = Max tpm limit for org
- `rpm_limit`: *Optional[int]* = Max rpm limit for org
- `model_max_budget`: *Optional[dict]* = Max budget for a specific model
- `budget_duration`: *Optional[str]* = Frequency of reseting org budget
Case 1: Create new org **without** a budget_id
```bash
curl --location 'http://0.0.0.0:4000/organization/new' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"organization_alias": "my-secret-org",
"models": ["model1", "model2"],
"max_budget": 100
}'
```
Case 2: Create new org **with** a budget_id
```bash
curl --location 'http://0.0.0.0:4000/organization/new' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"organization_alias": "my-secret-org",
"models": ["model1", "model2"],
"budget_id": "428eeaa8-f3ac-4e85-a8fb-7dc8d7aa8689"
}'
```
"""
global prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
if (
user_api_key_dict.user_role is None
or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
):
raise HTTPException(
status_code=401,
detail={
"error": f"Only admins can create orgs. Your role is = {user_api_key_dict.user_role}"
},
)
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
"""
Ensure only models that user has access to, are given to org
"""
if len(user_api_key_dict.models) == 0: # user has access to all models
pass
else:
if len(data.models) == 0:
raise HTTPException(
status_code=400,
detail={
"error": "User not allowed to give access to all models. Select models you want org to have access to."
},
)
for m in data.models:
if m not in user_api_key_dict.models:
raise HTTPException(
status_code=400,
detail={
"error": f"User not allowed to give access to model={m}. Models you have access to = {user_api_key_dict.models}"
},
)
organization_row = LiteLLM_OrganizationTable(
**data.json(exclude_none=True),
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,
)
new_organization_row = prisma_client.jsonify_object(
organization_row.json(exclude_none=True)
)
response = await prisma_client.db.litellm_organizationtable.create(
data={
**new_organization_row, # type: ignore
}
)
return response
@router.post(
"/organization/update",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_organization():
"""[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues"""
pass
@router.post(
"/organization/delete",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_organization():
"""[TODO] Not Implemented yet. Let us know if you need this - https://github.com/BerriAI/litellm/issues"""
pass
@router.post(
"/organization/info",
tags=["organization management"],
dependencies=[Depends(user_api_key_auth)],
)
async def info_organization(data: OrganizationRequest):
"""
Get the org specific information
"""
global prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
if len(data.organizations) == 0:
raise HTTPException(
status_code=400,
detail={
"error": f"Specify list of organization id's to query. Passed in={data.organizations}"
},
)
response = await prisma_client.db.litellm_organizationtable.find_many(
where={"organization_id": {"in": data.organizations}},
include={"litellm_budget_table": True},
)
return response
#### BUDGET TABLE MANAGEMENT ####
@ -8181,6 +7990,14 @@ async def login(request: Request):
# check if we can find the `username` in the db. on the ui, users can enter username=their email
_user_row = None
user_role: Optional[
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
] = None
if prisma_client is not None:
_user_row = await prisma_client.db.litellm_usertable.find_first(
where={"user_email": {"equals": username}}
@ -9654,6 +9471,7 @@ app.include_router(key_management_router)
app.include_router(internal_user_router)
app.include_router(team_router)
app.include_router(ui_sso_router)
app.include_router(organization_router)
app.include_router(spend_management_router)
app.include_router(caching_router)
app.include_router(analytics_router)

22
litellm/proxy/schema.prisma
View file

@ -26,6 +26,7 @@ model LiteLLM_BudgetTable {
keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
}
// Models on proxy
@ -118,7 +119,10 @@ model LiteLLM_UserTable {
allowed_cache_controls String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
// relations
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
organization_memberships LiteLLM_OrganizationMembership[]
invitations_created LiteLLM_InvitationLink[] @relation("CreatedBy")
invitations_updated LiteLLM_InvitationLink[] @relation("UpdatedBy")
invitations_user LiteLLM_InvitationLink[] @relation("UserId")
@ -232,6 +236,24 @@ model LiteLLM_TeamMembership {
@@id([user_id, team_id])
}
model LiteLLM_OrganizationMembership {
// Use this table to track Internal User and Organization membership. Helps tracking a users role within an Organization
user_id String?
organization_id String?
user_role String?
spend Float? @default(0.0)
budget_id String?
created_at DateTime? @default(now()) @map("created_at")
updated_at DateTime? @default(now()) @updatedAt @map("updated_at")
// relations
user LiteLLM_UserTable @relation(fields: [user_id], references: [user_id])
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
@@id([user_id, organization_id])
@@unique([user_id, organization_id])
}
model LiteLLM_InvitationLink {
// use this table to track invite links sent by admin for people to join the proxy
id String @id @default(uuid())

22
schema.prisma
View file

@ -26,6 +26,7 @@ model LiteLLM_BudgetTable {
keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
}
// Models on proxy
@ -118,7 +119,10 @@ model LiteLLM_UserTable {
allowed_cache_controls String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
// relations
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
organization_memberships LiteLLM_OrganizationMembership[]
invitations_created LiteLLM_InvitationLink[] @relation("CreatedBy")
invitations_updated LiteLLM_InvitationLink[] @relation("UpdatedBy")
invitations_user LiteLLM_InvitationLink[] @relation("UserId")
@ -232,6 +236,24 @@ model LiteLLM_TeamMembership {
@@id([user_id, team_id])
}
model LiteLLM_OrganizationMembership {
// Use this table to track Internal User and Organization membership. Helps tracking a users role within an Organization
user_id String?
organization_id String?
user_role String?
spend Float? @default(0.0)
budget_id String?
created_at DateTime? @default(now()) @map("created_at")
updated_at DateTime? @default(now()) @updatedAt @map("updated_at")
// relations
user LiteLLM_UserTable @relation(fields: [user_id], references: [user_id])
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
@@id([user_id, organization_id])
@@unique([user_id, organization_id])
}
model LiteLLM_InvitationLink {
// use this table to track invite links sent by admin for people to join the proxy
id String @id @default(uuid())

439
tests/proxy_admin_ui_tests/test_role_based_access.py Normal file
View file

@ -0,0 +1,439 @@
"""
RBAC tests
"""
import os
import sys
import traceback
import uuid
from datetime import datetime
from dotenv import load_dotenv
from fastapi import Request
from fastapi.routing import APIRoute
load_dotenv()
import io
import os
import time
# this file is to test litellm/proxy
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import logging
import pytest
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.auth_checks import get_user_object
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_key_fn,
generate_key_fn,
generate_key_helper_fn,
info_key_fn,
regenerate_key_fn,
update_key_fn,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.organization_endpoints import (
new_organization,
organization_member_add,
)
from litellm.proxy.management_endpoints.team_endpoints import (
new_team,
team_info,
update_team,
)
from litellm.proxy.proxy_server import (
LitellmUserRoles,
audio_transcriptions,
chat_completion,
completion,
embeddings,
image_generation,
model_list,
moderations,
new_end_user,
user_api_key_auth,
)
from litellm.proxy.spend_tracking.spend_management_endpoints import (
global_spend,
global_spend_logs,
global_spend_models,
global_spend_keys,
spend_key_fn,
spend_user_fn,
view_spend_logs,
)
from starlette.datastructures import URL
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend
verbose_proxy_logger.setLevel(level=logging.DEBUG)
from starlette.datastructures import URL
from litellm.caching import DualCache
from litellm.proxy._types import *
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
@pytest.fixture
def prisma_client():
from litellm.proxy.proxy_cli import append_query_params
### add connection pool + pool timeout args
params = {"connection_limit": 100, "pool_timeout": 60}
database_url = os.getenv("DATABASE_URL")
modified_url = append_query_params(database_url, params)
os.environ["DATABASE_URL"] = modified_url
# Assuming PrismaClient is a class that needs to be instantiated
prisma_client = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
# Reset litellm.proxy.proxy_server.prisma_client to None
litellm.proxy.proxy_server.litellm_proxy_budget_name = (
f"litellm-proxy-budget-{time.time()}"
)
litellm.proxy.proxy_server.user_custom_key_generate = None
return prisma_client
"""
RBAC Tests
1. Add a user to an organization
- test 1 - if organization_id does exist expect to create a new user and user, organization relation
2. org admin creates team in his org → success
3. org admin adds new internal user to his org → success
4. org admin creates team and internal user not in his org → fail both
"""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"user_role",
[
LitellmUserRoles.ORG_ADMIN,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
],
)
async def test_create_new_user_in_organization(prisma_client, user_role):
"""
Add a member to an organization and assert the user object is created with the correct organization memberships / roles
"""
master_key = "sk-1234"
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", master_key)
await litellm.proxy.proxy_server.prisma_client.connect()
created_user_id = f"new-user-{uuid.uuid4()}"
response = await new_organization(
data=NewOrganizationRequest(
organization_alias=f"new-org-{uuid.uuid4()}",
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
org_id = response.organization_id
response = await organization_member_add(
data=OrganizationMemberAddRequest(
organization_id=org_id,
member=Member(role=user_role, user_id=created_user_id),
),
http_request=None,
)
print("new user response", response)
# call get_user_object
user_object = await get_user_object(
user_id=created_user_id,
prisma_client=prisma_client,
user_api_key_cache=DualCache(),
user_id_upsert=False,
)
print("user object", user_object)
assert user_object.organization_memberships is not None
_membership = user_object.organization_memberships[0]
assert _membership.user_id == created_user_id
assert _membership.organization_id == org_id
if user_role != None:
assert _membership.user_role == user_role
else:
assert _membership.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
@pytest.mark.asyncio
async def test_org_admin_create_team_permissions(prisma_client):
"""
Create a new org admin
org admin creates a new team in their org -> success
"""
import json
master_key = "sk-1234"
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", master_key)
await litellm.proxy.proxy_server.prisma_client.connect()
response = await new_organization(
data=NewOrganizationRequest(
organization_alias=f"new-org-{uuid.uuid4()}",
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
org_id = response.organization_id
created_user_id = f"new-user-{uuid.uuid4()}"
response = await organization_member_add(
data=OrganizationMemberAddRequest(
organization_id=org_id,
member=Member(role=LitellmUserRoles.ORG_ADMIN, user_id=created_user_id),
),
http_request=None,
)
# create key with the response["user_id"]
# proxy admin will generate key for org admin
_new_key = await generate_key_fn(
data=GenerateKeyRequest(user_id=created_user_id),
user_api_key_dict=UserAPIKeyAuth(user_id=created_user_id),
)
new_key = _new_key.key
print("user api key auth response", response)
# Create /team/new request -> expect auth to pass
request = Request(scope={"type": "http"})
request._url = URL(url="/team/new")
async def return_body():
body = {"organization_id": org_id}
return bytes(json.dumps(body), "utf-8")
request.body = return_body
response = await user_api_key_auth(request=request, api_key="Bearer " + new_key)
# after auth - actually create team now
response = await new_team(
data=NewTeamRequest(
organization_id=org_id,
),
http_request=request,
user_api_key_dict=UserAPIKeyAuth(
user_id=response.user_id,
),
)
print("response from new team")
@pytest.mark.asyncio
async def test_org_admin_create_user_permissions(prisma_client):
"""
1. Create a new org admin
2. org admin adds a new member to their org -> success (using using /organization/member_add)
"""
import json
master_key = "sk-1234"
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", master_key)
await litellm.proxy.proxy_server.prisma_client.connect()
# create new org
response = await new_organization(
data=NewOrganizationRequest(
organization_alias=f"new-org-{uuid.uuid4()}",
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
# Create Org Admin
org_id = response.organization_id
created_user_id = f"new-user-{uuid.uuid4()}"
response = await organization_member_add(
data=OrganizationMemberAddRequest(
organization_id=org_id,
member=Member(role=LitellmUserRoles.ORG_ADMIN, user_id=created_user_id),
),
http_request=None,
)
# create key with for Org Admin
_new_key = await generate_key_fn(
data=GenerateKeyRequest(user_id=created_user_id),
user_api_key_dict=UserAPIKeyAuth(user_id=created_user_id),
)
new_key = _new_key.key
print("user api key auth response", response)
# Create /organization/member_add request -> expect auth to pass
request = Request(scope={"type": "http"})
request._url = URL(url="/organization/member_add")
async def return_body():
body = {"organization_id": org_id}
return bytes(json.dumps(body), "utf-8")
request.body = return_body
response = await user_api_key_auth(request=request, api_key="Bearer " + new_key)
# after auth - actually actually add new user to organization
new_internal_user_for_org = f"new-org-user-{uuid.uuid4()}"
response = await organization_member_add(
data=OrganizationMemberAddRequest(
organization_id=org_id,
member=Member(
role=LitellmUserRoles.INTERNAL_USER, user_id=new_internal_user_for_org
),
),
http_request=request,
)
print("response from new team")
@pytest.mark.asyncio
async def test_org_admin_create_user_team_wrong_org_permissions(prisma_client):
"""
Create a new org admin
org admin creates a new user and new team in orgs they are not part of -> expect error
"""
import json
master_key = "sk-1234"
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", master_key)
await litellm.proxy.proxy_server.prisma_client.connect()
created_user_id = f"new-user-{uuid.uuid4()}"
response = await new_organization(
data=NewOrganizationRequest(
organization_alias=f"new-org-{uuid.uuid4()}",
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
response2 = await new_organization(
data=NewOrganizationRequest(
organization_alias=f"new-org-{uuid.uuid4()}",
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
),
)
org1_id = response.organization_id # has an admin
org2_id = response2.organization_id # does not have an org admin
# Create Org Admin for Org1
created_user_id = f"new-user-{uuid.uuid4()}"
response = await organization_member_add(
data=OrganizationMemberAddRequest(
organization_id=org1_id,
member=Member(role=LitellmUserRoles.ORG_ADMIN, user_id=created_user_id),
),
http_request=None,
)
_new_key = await generate_key_fn(
data=GenerateKeyRequest(
user_id=created_user_id,
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.ORG_ADMIN,
user_id=created_user_id,
),
)
new_key = _new_key.key
print("user api key auth response", response)
# Add a new request in organization=org_without_admins -> expect fail (organization/member_add)
request = Request(scope={"type": "http"})
request._url = URL(url="/organization/member_add")
async def return_body():
body = {"organization_id": org2_id}
return bytes(json.dumps(body), "utf-8")
request.body = return_body
try:
response = await user_api_key_auth(request=request, api_key="Bearer " + new_key)
pytest.fail(
f"This should have failed!. creating a user in an org without admins"
)
except Exception as e:
print("got exception", e)
print("exception.message", e.message)
assert (
"You do not have a role within the selected organization. Passed organization_id"
in e.message
)
# Create /team/new request in organization=org_without_admins -> expect fail
request = Request(scope={"type": "http"})
request._url = URL(url="/team/new")
async def return_body():
body = {"organization_id": org2_id}
return bytes(json.dumps(body), "utf-8")
request.body = return_body
try:
response = await user_api_key_auth(request=request, api_key="Bearer " + new_key)
pytest.fail(
f"This should have failed!. Org Admin creating a team in an org where they are not an admin"
)
except Exception as e:
print("got exception", e)
print("exception.message", e.message)
assert (
"You do not have the required role to call" in e.message
and org2_id in e.message
)

2
tests/test_organizations.py
View file

@ -33,7 +33,7 @@ async def new_organization(session, i, organization_alias, max_budget=None):
@pytest.mark.asyncio
async def test_organization_new():
"""
Make 20 parallel calls to /user/new. Assert all worked.
Make 20 parallel calls to /organization/new. Assert all worked.
"""
organization_alias = f"Organization: {uuid.uuid4()}"
async with aiohttp.ClientSession() as session: