diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 5bb3e13ffe32..53b129782f55 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -271,7 +271,7 @@ def get_team_from_list( # response_model=UserInfoResponse, ) @management_endpoint_wrapper -async def user_info( # noqa: PLR0915 +async def user_info( user_id: Optional[str] = fastapi.Query( default=None, description="User ID in the request parameters" ), @@ -301,6 +301,11 @@ async def user_info( # noqa: PLR0915 raise Exception( "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 ## if user_id is not None: 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 ) # *NEW* get all teams in user 'teams' field - if ( - 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: + if caller_user_info is not None: teams_2 = await prisma_client.get_data( team_id_list=caller_user_info.teams, table_name="team", @@ -388,39 +381,7 @@ async def user_info( # noqa: PLR0915 user_info = {"spend": spend} ## REMOVE HASHED TOKEN INFO before returning ## - 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 = 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) + returned_keys = _process_keys_for_user_info(keys=keys, all_teams=teams_1) _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) +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: non_default_values = {} for k, v in data_json.items(): diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index e1fe3e4d3091..58cf22995939 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -735,3 +735,38 @@ def test_prepare_metadata_fields( updated_non_default_values = prepare_metadata_fields(**args) 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"