From 1014216d734d5dfcfb929d69bc2c9e8c49813040 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 22 Nov 2024 22:59:01 +0530 Subject: [PATCH] feat(key_management_endpoints.py): add support for restricting access to `/key/generate` by team/proxy level role Enables admin to restrict key creation, and assign team admins to handle distributing keys --- litellm/__init__.py | 2 + litellm/proxy/_new_secret_config.yaml | 6 +- litellm/proxy/_types.py | 4 - .../key_management_endpoints.py | 73 +++++++++++++++++++ litellm/types/utils.py | 13 ++++ 5 files changed, 93 insertions(+), 5 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index e6dc61dc7..65b1b3465 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -24,6 +24,7 @@ from litellm.proxy._types import ( KeyManagementSettings, LiteLLM_UpperboundKeyGenerateParams, ) +from litellm.types.utils import StandardKeyGenerationConfig import httpx import dotenv from enum import Enum @@ -273,6 +274,7 @@ s3_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None upperbound_key_generate_params: Optional[LiteLLM_UpperboundKeyGenerateParams] = None +key_generation_settings: Optional[StandardKeyGenerationConfig] = None default_internal_user_params: Optional[Dict] = None default_team_settings: Optional[List] = None max_user_budget: Optional[float] = None diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index f12226736..dd4c06576 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -15,4 +15,8 @@ model_list: litellm_settings: success_callback: ["langfuse"] callbacks: ["prometheus"] - # disable_end_user_cost_tracking: true + key_generation_settings: + team_key_generation: + allowed_team_member_roles: ["admin"] + personal_key_generation: # maps to 'Default Team' on UI + allowed_user_roles: ["proxy_admin"] \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b0d272f26..9e05e4cff 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -892,10 +892,6 @@ class DeleteCustomerRequest(LiteLLMBase): class Member(LiteLLMBase): role: Literal[ - LitellmUserRoles.ORG_ADMIN, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - # older Member roles "admin", "user", ] diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e4493a28c..ab13616d5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -40,6 +40,77 @@ from litellm.proxy.utils import ( ) from litellm.secret_managers.main import get_secret + +def _is_team_key(data: GenerateKeyRequest): + return data.team_id is not None + + +def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth): + if ( + litellm.key_generation_settings is None + or litellm.key_generation_settings.get("team_key_generation") is None + ): + return True + + if user_api_key_dict.team_member is None: + raise HTTPException( + status_code=400, + detail=f"User not assigned to team. Got team_member={user_api_key_dict.team_member}", + ) + + team_member_role = user_api_key_dict.team_member.role + if ( + team_member_role + not in litellm.key_generation_settings["team_key_generation"][ # type: ignore + "allowed_team_member_roles" + ] + ): + raise HTTPException( + status_code=400, + detail=f"Team member role {team_member_role} not in allowed_team_member_roles={litellm.key_generation_settings['team_key_generation']['allowed_team_member_roles']}", # type: ignore + ) + return True + + +def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth): + + if ( + litellm.key_generation_settings is None + or litellm.key_generation_settings.get("personal_key_generation") is None + ): + return True + + if ( + user_api_key_dict.user_role + not in litellm.key_generation_settings["personal_key_generation"][ # type: ignore + "allowed_user_roles" + ] + ): + raise HTTPException( + status_code=400, + detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}", # type: ignore + ) + return True + + +def key_generation_check( + user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest +) -> bool: + """ + Check if admin has restricted key creation to certain roles for teams or individuals + """ + if litellm.key_generation_settings is None: + return True + + ## check if key is for team or individual + is_team_key = _is_team_key(data=data) + + if is_team_key: + return _team_key_generation_check(user_api_key_dict) + else: + return _personal_key_generation_check(user_api_key_dict=user_api_key_dict) + + router = APIRouter() @@ -131,6 +202,8 @@ async def generate_key_fn( # noqa: PLR0915 raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=message ) + elif litellm.key_generation_settings is not None: + key_generation_check(user_api_key_dict=user_api_key_dict, data=data) # check if user set default key/generate params on config.yaml if litellm.default_key_generate_params is not None: for elem in data: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d02129681..334894320 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1602,3 +1602,16 @@ class StandardCallbackDynamicParams(TypedDict, total=False): langsmith_api_key: Optional[str] langsmith_project: Optional[str] langsmith_base_url: Optional[str] + + +class TeamUIKeyGenerationConfig(TypedDict): + allowed_team_member_roles: List[str] + + +class PersonalUIKeyGenerationConfig(TypedDict): + allowed_user_roles: List[str] + + +class StandardKeyGenerationConfig(TypedDict, total=False): + team_key_generation: TeamUIKeyGenerationConfig + personal_key_generation: PersonalUIKeyGenerationConfig