diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e535ccaaa4..5a0a05114a 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -614,7 +614,7 @@ async def user_info( user_id is None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ): - return await _get_user_info_for_proxy_admin() + return await _get_user_info_for_proxy_admin(user_api_key_dict=user_api_key_dict) elif user_id is None: user_id = user_api_key_dict.user_id ## GET USER ROW ## @@ -714,7 +714,7 @@ async def user_info( raise handle_exception_on_proxy(e) -async def _get_user_info_for_proxy_admin(): +async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): """ Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying @@ -754,9 +754,23 @@ async def _get_user_info_for_proxy_admin(): _teams_in_db = [LiteLLM_TeamTable(**team) for team in _teams_in_db] _teams_in_db.sort(key=lambda x: (getattr(x, "team_alias", "") or "")) returned_keys = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) + + # Get admin's own user_id and user_info + admin_user_id = user_api_key_dict.user_id + admin_user_info = None + + if admin_user_id is not None: + admin_user_info = await prisma_client.get_data(user_id=admin_user_id) + if admin_user_info is not None: + admin_user_info = ( + admin_user_info.model_dump() + if isinstance(admin_user_info, BaseModel) + else admin_user_info + ) + return UserInfoResponse( - user_id=None, - user_info=None, + user_id=admin_user_id, + user_info=admin_user_info, keys=returned_keys, teams=_teams_in_db, ) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index cc1cc27839..fd97a38b41 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2,7 +2,8 @@ import asyncio import json import os import sys -from typing import Any, Dict, List, Optional +from datetime import datetime +from typing import Any, Dict, List, Optional, Union from unittest.mock import Mock import pytest @@ -1486,12 +1487,46 @@ class MockPrismaClientDB: mock_key_data, ): self.db = MockDb(mock_team_data, mock_key_data) + + async def get_data( + self, + token: Optional[Union[str, list]] = None, + user_id: Optional[str] = None, + user_id_list: Optional[list] = None, + team_id: Optional[str] = None, + team_id_list: Optional[list] = None, + key_val: Optional[dict] = None, + table_name: Optional[str] = None, + query_type: str = "find_unique", + expires: Optional[datetime] = None, + reset_at: Optional[datetime] = None, + offset: Optional[int] = None, + limit: Optional[int] = None, + ): + """Mock get_data method to return user info for admin""" + from litellm.proxy._types import LiteLLM_UserTable + + # Return a proper LiteLLM_UserTable object when querying by user_id + if user_id: + return LiteLLM_UserTable( + user_id=user_id, + user_role="proxy_admin", + spend=0.0, + max_budget=None, + ) + return None @pytest.mark.asyncio async def test_get_user_info_for_proxy_admin(mock_team_data, mock_key_data): # Patch the prisma_client import - from litellm.proxy._types import UserInfoResponse + from litellm.proxy._types import UserAPIKeyAuth, UserInfoResponse + + # Create a mock user_api_key_dict for admin user + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user_123", + user_role="proxy_admin", + ) with patch( "litellm.proxy.proxy_server.prisma_client", @@ -1502,11 +1537,18 @@ async def test_get_user_info_for_proxy_admin(mock_team_data, mock_key_data): ) # Execute the function - result = await _get_user_info_for_proxy_admin() + result = await _get_user_info_for_proxy_admin( + user_api_key_dict=mock_user_api_key_dict + ) # Verify the result structure assert isinstance(result, UserInfoResponse) assert len(result.keys) == 2 + # Verify admin's user_id is populated + assert result.user_id == "admin_user_123" + # Verify admin's user_info is populated + assert result.user_info is not None + assert result.user_info["user_id"] == "admin_user_123" def test_custom_openid_response():