forked from phoenix/litellm-mirror
(feat proxy) [beta] add support for organization role based access controls (#6112)
* 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
GitHub
parent
945267a511
commit
1fd437e263
14 changed files with 1474 additions and 261 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue