From de551889f3459b376946ae46aa948dee2d65c42b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 14:20:04 -0800 Subject: [PATCH 1/2] Add optional field expand to /key/list --- litellm/proxy/_types.py | 1 + .../key_management_endpoints.py | 26 +++- .../test_key_management_endpoints.py | 112 ++++++++++++++++++ 3 files changed, 138 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 570a6bb9f3..a94b4d4a07 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2152,6 +2152,7 @@ class UserAPIKeyAuth( user_rpm_limit: Optional[int] = None user_email: Optional[str] = None request_route: Optional[str] = None + user: Optional[Any] = None # Expanded user object when expand=user is used model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index cc2ac90814..65b4ec11f4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3020,10 +3020,14 @@ async def list_keys( description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), + expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. + Parameters: + expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) + Returns: { "keys": List[str] or List[UserAPIKeyAuth], @@ -3031,6 +3035,9 @@ async def list_keys( "current_page": int, "total_pages": int, } + + When expand includes "user", each key object will include a "user" field with the associated user object. + Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter. """ try: from litellm.proxy.proxy_server import prisma_client @@ -3080,6 +3087,7 @@ async def list_keys( include_created_by_keys=include_created_by_keys, sort_by=sort_by, sort_order=sort_order, + expand=expand, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -3232,6 +3240,7 @@ async def _list_key_helper( include_created_by_keys: bool = False, sort_by: Optional[str] = None, sort_order: str = "desc", + expand: Optional[List[str]] = None, ) -> KeyListResponseObject: """ Helper function to list keys @@ -3334,13 +3343,28 @@ async def _list_key_helper( # Calculate total pages total_pages = -(-total_count // size) # Ceiling division + # Fetch user information if expand includes "user" + user_map = {} + if expand and "user" in expand: + user_ids = [key.user_id for key in keys if key.user_id] + if user_ids: + users = await prisma_client.db.litellm_usertable.find_many( + where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates + ) + user_map = {user.user_id: user for user in users} + # Prepare response key_list: List[Union[str, UserAPIKeyAuth]] = [] for key in keys: key_dict = key.dict() # Attach object_permission if object_permission_id is set key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) - if return_full_object is True: + + # Include user information if expand includes "user" + if expand and "user" in expand and key.user_id and key.user_id in user_map: + key_dict["user"] = user_map[key.user_id].dict() + + if return_full_object is True or (expand and "user" in expand): key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object else: _token = key_dict.get("token") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 648045a7ea..cba06e7fb3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3405,3 +3405,115 @@ async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch ) assert result is False + + +@pytest.mark.asyncio +async def test_list_keys_with_expand_user(): + """ + Test that expand=user parameter correctly includes user information in the response. + """ + mock_prisma_client = AsyncMock() + + # Create mock keys with user_ids + mock_key1 = MagicMock() + mock_key1.token = "token1" + mock_key1.user_id = "user123" + mock_key1.dict.return_value = { + "token": "token1", + "user_id": "user123", + "key_alias": "key1", + "models": ["gpt-4"], + } + + mock_key2 = MagicMock() + mock_key2.token = "token2" + mock_key2.user_id = "user456" + mock_key2.dict.return_value = { + "token": "token2", + "user_id": "user456", + "key_alias": "key2", + "models": ["gpt-3.5-turbo"], + } + + mock_find_many_keys = AsyncMock(return_value=[mock_key1, mock_key2]) + mock_count_keys = AsyncMock(return_value=2) + + # Create mock users + mock_user1 = MagicMock() + mock_user1.user_id = "user123" + mock_user1.user_email = "user1@example.com" + mock_user1.dict.return_value = { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + + mock_user2 = MagicMock() + mock_user2.user_id = "user456" + mock_user2.user_email = "user2@example.com" + mock_user2.dict.return_value = { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } + + mock_find_many_users = AsyncMock(return_value=[mock_user1, mock_user2]) + + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys + mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users + + args = { + "prisma_client": mock_prisma_client, + "page": 1, + "size": 50, + "user_id": None, + "team_id": None, + "organization_id": None, + "key_alias": None, + "key_hash": None, + "exclude_team_id": None, + "return_full_object": False, # This should be overridden by expand=user + "admin_team_ids": None, + "include_created_by_keys": False, + "expand": ["user"], # Test the expand parameter + } + + result = await _list_key_helper(**args) + + # Verify that keys were fetched + mock_find_many_keys.assert_called_once() + mock_count_keys.assert_called_once() + + # Verify that users were fetched + # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present + call_args = mock_find_many_users.call_args + assert call_args is not None + where_clause = call_args.kwargs["where"] + assert "user_id" in where_clause + assert "in" in where_clause["user_id"] + user_ids_in_query = set(where_clause["user_id"]["in"]) + assert user_ids_in_query == {"user123", "user456"} + + # Verify response structure + assert len(result["keys"]) == 2 + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 + + # Verify that user data is included in the response + # Since expand=user is specified, keys should be full objects + assert isinstance(result["keys"][0], UserAPIKeyAuth) + assert isinstance(result["keys"][1], UserAPIKeyAuth) + + # Verify user data is attached to keys + assert result["keys"][0].user == { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + assert result["keys"][1].user == { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } From 03cee685881025be13584ca4729725b8985dde1c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 29 Dec 2025 14:27:55 -0800 Subject: [PATCH 2/2] Ruff check --- .../key_management_endpoints.py | 94 ++++++++++++------- 1 file changed, 59 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 65b4ec11f4..4f9534521c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3223,46 +3223,17 @@ def _validate_sort_params( return order_by -async def _list_key_helper( - prisma_client: PrismaClient, - page: int, - size: int, +def _build_key_filter_conditions( user_id: Optional[str], team_id: Optional[str], organization_id: Optional[str], key_alias: Optional[str], key_hash: Optional[str], - exclude_team_id: Optional[str] = None, - return_full_object: bool = False, - admin_team_ids: Optional[ - List[str] - ] = None, # New parameter for teams where user is admin - include_created_by_keys: bool = False, - sort_by: Optional[str] = None, - sort_order: str = "desc", - expand: Optional[List[str]] = None, -) -> KeyListResponseObject: - """ - Helper function to list keys - Args: - page: int - size: int - user_id: Optional[str] - team_id: Optional[str] - key_alias: Optional[str] - exclude_team_id: Optional[str] # exclude a specific team_id - return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token - admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin - - Returns: - KeyListResponseObject - { - "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types - "total_count": int, - "current_page": int, - "total_pages": int, - } - """ + exclude_team_id: Optional[str], + admin_team_ids: Optional[List[str]], + include_created_by_keys: bool, +) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: + """Build filter conditions for key listing.""" # Prepare filter conditions where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) @@ -3303,6 +3274,59 @@ async def _list_key_helper( where.update(or_conditions[0]) verbose_proxy_logger.debug(f"Filter conditions: {where}") + return where + + +async def _list_key_helper( + prisma_client: PrismaClient, + page: int, + size: int, + user_id: Optional[str], + team_id: Optional[str], + organization_id: Optional[str], + key_alias: Optional[str], + key_hash: Optional[str], + exclude_team_id: Optional[str] = None, + return_full_object: bool = False, + admin_team_ids: Optional[ + List[str] + ] = None, # New parameter for teams where user is admin + include_created_by_keys: bool = False, + sort_by: Optional[str] = None, + sort_order: str = "desc", + expand: Optional[List[str]] = None, +) -> KeyListResponseObject: + """ + Helper function to list keys + Args: + page: int + size: int + user_id: Optional[str] + team_id: Optional[str] + key_alias: Optional[str] + exclude_team_id: Optional[str] # exclude a specific team_id + return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token + admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin + + Returns: + KeyListResponseObject + { + "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types + "total_count": int, + "current_page": int, + "total_pages": int, + } + """ + where = _build_key_filter_conditions( + user_id=user_id, + team_id=team_id, + organization_id=organization_id, + key_alias=key_alias, + key_hash=key_hash, + exclude_team_id=exclude_team_id, + admin_team_ids=admin_team_ids, + include_created_by_keys=include_created_by_keys, + ) # Calculate skip for pagination skip = (page - 1) * size