Add optional field expand to /key/list

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)
 

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")

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",
+    }