(fix) litellm bug in

This commit is contained in:
ishaan-jaff 2024-02-17 18:41:37 -08:00
parent 5bb3d346bb
commit fdf3cd4c59
2 changed files with 34 additions and 14 deletions

View file

@ -93,6 +93,7 @@ from litellm.proxy.utils import (
html_form, html_form,
_read_request_body, _read_request_body,
_is_valid_team_configs, _is_valid_team_configs,
_is_user_proxy_admin,
) )
from litellm.proxy.secret_managers.google_kms import load_google_kms from litellm.proxy.secret_managers.google_kms import load_google_kms
import pydantic import pydantic
@ -503,11 +504,7 @@ async def user_api_key_auth(
continue continue
assert isinstance(_user, dict) assert isinstance(_user, dict)
# check if user is admin # # check if user is admin #
if (
_user.get("user_role", None) is not None
and _user.get("user_role") == "proxy_admin"
):
return UserAPIKeyAuth(api_key=master_key)
# Token exists, not expired now check if its in budget for the user # Token exists, not expired now check if its in budget for the user
user_max_budget = _user.get("max_budget", None) user_max_budget = _user.get("max_budget", None)
user_current_spend = _user.get("spend", None) user_current_spend = _user.get("spend", None)
@ -642,11 +639,15 @@ async def user_api_key_auth(
) )
) )
if ( if (
(
route.startswith("/key/") route.startswith("/key/")
or route.startswith("/user/") or route.startswith("/user/")
or route.startswith("/model/") or route.startswith("/model/")
or route.startswith("/spend/") or route.startswith("/spend/")
) and (not is_master_key_valid): )
and (not is_master_key_valid)
and (not _is_user_proxy_admin(user_id_information))
):
allow_user_auth = False allow_user_auth = False
if ( if (
general_settings.get("allow_user_auth", False) == True general_settings.get("allow_user_auth", False) == True
@ -737,6 +738,9 @@ async def user_api_key_auth(
): ):
# Do something if the current route starts with any of the allowed routes # Do something if the current route starts with any of the allowed routes
pass pass
else:
if _is_user_proxy_admin(user_id_information):
pass
else: else:
raise Exception( raise Exception(
f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed" f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
@ -4944,7 +4948,7 @@ async def auth_callback(request: Request):
if user_id is None: if user_id is None:
user_id = getattr(result, "first_name", "") + getattr(result, "last_name", "") user_id = getattr(result, "first_name", "") + getattr(result, "last_name", "")
response = await generate_key_helper_fn( response = await generate_key_helper_fn(
**{"duration": "1hr", "key_max_budget": 0, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id, "team_id": "litellm-dashboard", "user_email": user_email} # type: ignore **{"duration": "1hr", "key_max_budget": 0.01, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id, "team_id": "litellm-dashboard", "user_email": user_email} # type: ignore
) )
key = response["token"] # type: ignore key = response["token"] # type: ignore
user_id = response["user_id"] # type: ignore user_id = response["user_id"] # type: ignore

View file

@ -1408,6 +1408,22 @@ def _is_valid_team_configs(team_id=None, team_config=None, request_data=None):
return return
def _is_user_proxy_admin(user_id_information=None):
if (
user_id_information == None
or len(user_id_information) == 0
or user_id_information[0] == None
):
return False
_user = user_id_information[0]
if (
_user.get("user_role", None) is not None
and _user.get("user_role") == "proxy_admin"
):
return True
return False
# LiteLLM Admin UI - Non SSO Login # LiteLLM Admin UI - Non SSO Login
html_form = """ html_form = """
<!DOCTYPE html> <!DOCTYPE html>