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): fix _get_user_role

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-05 20:54:06 -07:00
parent 2f304c92e4
commit 818a70838d
1 changed files with 21 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1171,32 +1171,13 @@ 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 = _get_user_role(
user_id_information=user_id_information
return _return_user_api_key_auth_obj(
user_id_information=user_id_information,
api_key=api_key,
parent_otel_span=parent_otel_span,
valid_token_dict=valid_token_dict,
route=route,
)
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,
)
else:
raise Exception()
except Exception as e:
@ -1244,7 +1225,10 @@ def _return_user_api_key_auth_obj(
valid_token_dict: dict,
route: str,
) -> UserAPIKeyAuth:
retrieved_user_role = _get_user_role(user_id_information=user_id_information)
retrieved_user_role = (
_get_user_role(user_id_information=user_id_information)
or LitellmUserRoles.INTERNAL_USER
)
if user_id_information is not None and _is_user_proxy_admin(user_id_information):
return UserAPIKeyAuth(
api_key=api_key,
@ -1318,19 +1302,21 @@ def _is_user_proxy_admin(user_id_information: Optional[list]):
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,
) -> Optional[
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 LitellmUserRoles.INTERNAL_USER
return None
if len(user_id_information) == 0 or user_id_information[0] is None:
return LitellmUserRoles.INTERNAL_USER
return None
_user = user_id_information[0]