phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(user_api_key_auth.py): handle older user_role's

Browse source 
Fixes issue where older user_role's (e.g. app_user) weren't being recognized. + Adds testing for it
This commit is contained in:
Krrish Dholakia 2024-08-05 08:57:06 -07:00
parent 51211dc2a6
commit f610fba58f
3 changed files with 77 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1164,7 +1164,7 @@ async def user_api_key_auth(
# No token was found when looking up in the DB
raise Exception("Invalid proxy server token passed")
if valid_token_dict is not None:
retrieved_user_role: Optional[str] = _get_user_role(
retrieved_user_role = _get_user_role(
user_id_information=user_id_information
)
if user_id_information is not None and _is_user_proxy_admin(
@ -1179,14 +1179,14 @@ async def user_api_key_auth(
elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value:
return UserAPIKeyAuth(
api_key=api_key,
user_role=retrieved_user_role or LitellmUserRoles.INTERNAL_USER, # type: ignore
user_role=retrieved_user_role,
parent_otel_span=parent_otel_span,
**valid_token_dict,
)
else:
return UserAPIKeyAuth(
api_key=api_key,
user_role=retrieved_user_role or LitellmUserRoles.INTERNAL_USER, # type: ignore
user_role=retrieved_user_role,
parent_otel_span=parent_otel_span,
**valid_token_dict,
)
@ -1230,6 +1230,37 @@ async def user_api_key_auth(
)
def _return_user_api_key_auth_obj(
user_id_information: Optional[list],
api_key: str,
parent_otel_span: Optional[Span],
valid_token_dict: dict,
route: str,
) -> UserAPIKeyAuth:
retrieved_user_role = _get_user_role(user_id_information=user_id_information)
if user_id_information is not None and _is_user_proxy_admin(user_id_information):
return UserAPIKeyAuth(
api_key=api_key,
user_role=LitellmUserRoles.PROXY_ADMIN,
parent_otel_span=parent_otel_span,
**valid_token_dict,
)
elif _has_user_setup_sso() and route in LiteLLMRoutes.sso_only_routes.value:
return UserAPIKeyAuth(
api_key=api_key,
user_role=retrieved_user_role,
parent_otel_span=parent_otel_span,
**valid_token_dict,
)
else:
return UserAPIKeyAuth(
api_key=api_key,
user_role=retrieved_user_role,
parent_otel_span=parent_otel_span,
**valid_token_dict,
)
def _has_user_setup_sso():
"""
Check if the user has set up single sign-on (SSO) by verifying the presence of Microsoft client ID, Google client ID, and UI username environment variables.
@ -1278,15 +1309,28 @@ def _is_user_proxy_admin(user_id_information: Optional[list]):
return False
def _get_user_role(user_id_information: Optional[list]) -> Optional[str]:
def _get_user_role(
user_id_information: Optional[list],
) -> Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.TEAM,
LitellmUserRoles.CUSTOMER,
]:
if user_id_information is None:
return None
return LitellmUserRoles.INTERNAL_USER
if len(user_id_information) == 0 or user_id_information[0] is None:
return None
return LitellmUserRoles.INTERNAL_USER
_user = user_id_information[0]
return _user.get("user_role")
_user_role = _user.get("user_role")
if _user_role in list(LitellmUserRoles.__annotations__.keys()):
return _user_role
return LitellmUserRoles.INTERNAL_USER
def _check_valid_ip(allowed_ips: Optional[List[str]], request: Request) -> bool: