(UI) Load time improvement - Sub 2s load time for Home Page ️ (#7014)

* ui fix LiteLLM_VerificationToken

* speed up ui load time for proxy admin

* undo type change

* fix _get_user_info_for_proxy_admin

* test_user_info_as_proxy_admin

* fix linting error

* fix merge conflicts
This commit is contained in:
Ishaan Jaff 2024-12-04 17:55:26 -08:00 committed by GitHub
parent 8c4b1de69f
commit 0b483078c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 130 additions and 47 deletions

View file

@ -271,7 +271,7 @@ def get_team_from_list(
# response_model=UserInfoResponse, # response_model=UserInfoResponse,
) )
@management_endpoint_wrapper @management_endpoint_wrapper
async def user_info( # noqa: PLR0915 async def user_info(
user_id: Optional[str] = fastapi.Query( user_id: Optional[str] = fastapi.Query(
default=None, description="User ID in the request parameters" default=None, description="User ID in the request parameters"
), ),
@ -301,6 +301,11 @@ async def user_info( # noqa: PLR0915
raise Exception( raise Exception(
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
) )
if (
user_id is None
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
):
return await _get_user_info_for_proxy_admin()
## GET USER ROW ## ## GET USER ROW ##
if user_id is not None: if user_id is not None:
user_info = await prisma_client.get_data(user_id=user_id) user_info = await prisma_client.get_data(user_id=user_id)
@ -348,19 +353,7 @@ async def user_info( # noqa: PLR0915
user_id=user_api_key_dict.user_id user_id=user_api_key_dict.user_id
) )
# *NEW* get all teams in user 'teams' field # *NEW* get all teams in user 'teams' field
if ( if caller_user_info is not None:
getattr(caller_user_info, "user_role", None)
== LitellmUserRoles.PROXY_ADMIN
):
from litellm.proxy.management_endpoints.team_endpoints import list_team
teams_2 = await list_team(
http_request=Request(
scope={"type": "http", "path": "/user/info"},
),
user_api_key_dict=user_api_key_dict,
)
elif caller_user_info is not None:
teams_2 = await prisma_client.get_data( teams_2 = await prisma_client.get_data(
team_id_list=caller_user_info.teams, team_id_list=caller_user_info.teams,
table_name="team", table_name="team",
@ -388,39 +381,7 @@ async def user_info( # noqa: PLR0915
user_info = {"spend": spend} user_info = {"spend": spend}
## REMOVE HASHED TOKEN INFO before returning ## ## REMOVE HASHED TOKEN INFO before returning ##
returned_keys = [] returned_keys = _process_keys_for_user_info(keys=keys, all_teams=teams_1)
if keys is None:
pass
else:
for key in keys:
if (
key.token == litellm_master_key_hash
and general_settings.get("disable_master_key_return", False)
is True ## [IMPORTANT] used by hosted proxy-ui to prevent sharing master key on ui
):
continue
try:
key = key.model_dump() # noqa
except Exception:
# if using pydantic v1
key = key.dict()
if (
"team_id" in key
and key["team_id"] is not None
and key["team_id"] != "litellm-dashboard"
):
team_info = get_team_from_list(
team_list=teams_1, team_id=key["team_id"]
)
if team_info is not None:
team_alias = getattr(team_info, "team_alias", None)
key["team_alias"] = team_alias
else:
key["team_alias"] = None
else:
key["team_alias"] = "None"
returned_keys.append(key)
_user_info = ( _user_info = (
user_info.model_dump() if isinstance(user_info, BaseModel) else user_info user_info.model_dump() if isinstance(user_info, BaseModel) else user_info
@ -439,6 +400,93 @@ async def user_info( # noqa: PLR0915
raise handle_exception_on_proxy(e) raise handle_exception_on_proxy(e)
async def _get_user_info_for_proxy_admin():
"""
Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying
- get all teams in LiteLLM_TeamTable
- get all keys in LiteLLM_VerificationToken table
Why separate helper for proxy admin ?
- To get Faster UI load times, get all teams and virtual keys in 1 query
"""
from litellm.proxy.proxy_server import prisma_client
sql_query = """
SELECT
(SELECT json_agg(t.*) FROM "LiteLLM_TeamTable" t) as teams,
(SELECT json_agg(k.*) FROM "LiteLLM_VerificationToken" k WHERE k.team_id != 'litellm-dashboard') as keys
"""
if prisma_client is None:
raise Exception(
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
results = await prisma_client.db.query_raw(sql_query)
_keys_in_db = results[0]["keys"]
# cast all keys to LiteLLM_VerificationToken
keys_in_db = []
for key in _keys_in_db:
if key.get("models") is None:
key["models"] = []
keys_in_db.append(LiteLLM_VerificationToken(**key))
# cast all teams to LiteLLM_TeamTable
_teams_in_db = results[0]["teams"]
_teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db]
returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db)
return UserInfoResponse(
user_id=None,
user_info=None,
keys=returned_keys,
teams=_teams_in_db,
)
def _process_keys_for_user_info(
keys: Optional[List[LiteLLM_VerificationToken]],
all_teams: Optional[Union[List[LiteLLM_TeamTable], List[TeamListResponseObject]]],
):
from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash
returned_keys = []
if keys is None:
pass
else:
for key in keys:
if (
key.token == litellm_master_key_hash
and general_settings.get("disable_master_key_return", False)
is True ## [IMPORTANT] used by hosted proxy-ui to prevent sharing master key on ui
):
continue
try:
_key: dict = key.model_dump() # noqa
except Exception:
# if using pydantic v1
_key = key.dict()
if (
"team_id" in _key
and _key["team_id"] is not None
and _key["team_id"] != "litellm-dashboard"
):
team_info = get_team_from_list(
team_list=all_teams, team_id=_key["team_id"]
)
if team_info is not None:
team_alias = getattr(team_info, "team_alias", None)
_key["team_alias"] = team_alias
else:
_key["team_alias"] = None
else:
_key["team_alias"] = "None"
returned_keys.append(_key)
return returned_keys
def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> dict: def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> dict:
non_default_values = {} non_default_values = {}
for k, v in data_json.items(): for k, v in data_json.items():

View file

@ -735,3 +735,38 @@ def test_prepare_metadata_fields(
updated_non_default_values = prepare_metadata_fields(**args) updated_non_default_values = prepare_metadata_fields(**args)
assert updated_non_default_values == expected_result assert updated_non_default_values == expected_result
@pytest.mark.asyncio
async def test_user_info_as_proxy_admin(prisma_client):
"""
Test /user/info endpoint as a proxy admin without passing a user ID.
Verifies that the endpoint returns all teams and keys.
"""
litellm.set_verbose = True
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Call user_info as a proxy admin without a user_id
user_info_response = await user_info(
user_id=None,
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="admin",
),
)
print("user info response: ", user_info_response.model_dump_json(indent=4))
# Verify response
assert user_info_response.user_id is None
assert user_info_response.user_info is None
# Verify that teams and keys are returned
assert user_info_response.teams is not None
assert len(user_info_response.teams) > 0, "Expected at least one team in response"
assert user_info_response.keys is not None
assert len(user_info_response.keys) > 0, "Expected at least one key in response"