Merge pull request #18502 from BerriAI/litellm_key_list_user_object

[Feature] Add Optional Query Param "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")
@@ -3215,45 +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",
-) -> 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())
@@ -3294,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
@@ -3334,13 +3367,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",
+    }