diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b07d44deb6..947274329cd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -481,6 +481,7 @@ class LiteLLMRoutes(enum.Enum): "/organization/list", "/team/available", "/user/info", + "/v2/user/info", "/model/info", "/v1/model/info", "/v2/model/info", diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index a00401008fc..68631e211cb 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -170,7 +170,7 @@ class RouteChecks: if route == "/key/info": # handled by function itself pass - elif route == "/user/info": + elif route in ("/user/info", "/v2/user/info"): # check if user can access this route query_params = request.query_params user_id = query_params.get("user_id") diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 5a0a05114a3..5ab57ba67ba 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -19,6 +19,7 @@ from typing import Any, Dict, List, Optional, Union, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +from fastapi.responses import JSONResponse import litellm from litellm._logging import verbose_proxy_logger @@ -43,6 +44,7 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, + UserInfoV2Response, UserListResponse, UserUpdateResult, ) @@ -704,7 +706,12 @@ async def user_info( user_id=user_id, user_info=_user_info, keys=returned_keys, teams=team_list ) - return response_data + response = JSONResponse(content=response_data.model_dump(mode="json")) + response.headers["Deprecation"] = "true" + response.headers[ + "Link" + ] = '; rel="successor-version"' + return response except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.user_info(): Exception occured - {}".format( @@ -714,6 +721,73 @@ async def user_info( raise handle_exception_on_proxy(e) +@router.get( + "/v2/user/info", + tags=["Internal User management"], + dependencies=[Depends(user_api_key_auth)], + response_model=UserInfoV2Response, +) +@management_endpoint_wrapper +async def user_info_v2( + request: Request, + user_id: Optional[str] = fastapi.Query( + default=None, description="User ID in the request parameters" + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Lightweight endpoint that returns only the user profile. + + Use /key/list?user_id= for keys and + /v2/team/list?user_id= for teams. + + Example request: + ``` + curl -X GET 'http://localhost:4000/v2/user/info?user_id=krrish7%40berri.ai' \ + --header 'Authorization: Bearer sk-1234' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + try: + if user_id is not None and " " in user_id: + user_id = get_user_id_from_request(request=request) + + 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" + ) + + if user_id is None: + user_id = user_api_key_dict.user_id + + if user_id is None: + raise HTTPException( + status_code=400, + detail="user_id is required", + ) + + user_info = await prisma_client.get_data(user_id=user_id) + + if user_info is None: + raise HTTPException( + status_code=404, + detail=f"User {user_id} not found", + ) + + return UserInfoV2Response( + user_id=user_id, + user_info=user_info, + ) + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.user_info_v2(): Exception occured - {}".format( + str(e) + ) + ) + raise handle_exception_on_proxy(e) + + 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 diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 6023094a920..01a89f53a1f 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -4,6 +4,7 @@ from fastapi import HTTPException from pydantic import BaseModel, EmailStr, field_validator from litellm.proxy._types import ( + LiteLLM_UserTable, LiteLLM_UserTableWithKeyCount, UpdateUserRequest, UpdateUserRequestNoUserIDorEmail, @@ -22,6 +23,18 @@ class UserListResponse(BaseModel): total_pages: int +class UserInfoV2Response(BaseModel): + """ + Response model for the /v2/user/info endpoint. + + Returns only the user profile. Keys and teams should be fetched + via /key/list and /v2/team/list respectively. + """ + + user_id: str + user_info: LiteLLM_UserTable + + class BulkUpdateUserRequest(BaseModel): """Request for bulk user updates""" diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index f1e96f3e660..cb5f86c07e2 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1190,3 +1190,62 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re request_data={}, ) assert "Only proxy admin can be used to generate" in str(exc_info.value) + + +def test_v2_user_info_non_admin_blocked_for_other_user(): + """ + Test that /v2/user/info blocks non-admin users from querying another user's info. + """ + user_obj = LiteLLM_UserTable( + user_id="user-A", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + valid_token = UserAPIKeyAuth( + user_id="user-A", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + request = MagicMock(spec=Request) + request.query_params = {"user_id": "user-B"} + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/v2/user/info", + request=request, + valid_token=valid_token, + request_data={}, + ) + + assert exc_info.value.status_code == 403 + assert "key not allowed to access this user's info" in str(exc_info.value.detail) + + +def test_v2_user_info_non_admin_allowed_for_own_user(): + """ + Test that /v2/user/info allows non-admin users to query their own info. + """ + user_obj = LiteLLM_UserTable( + user_id="user-A", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + valid_token = UserAPIKeyAuth( + user_id="user-A", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + request = MagicMock(spec=Request) + request.query_params = {"user_id": "user-A"} + + # Should not raise + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/v2/user/info", + request=request, + valid_token=valid_token, + request_data={}, + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 839885bc752..cf99b0da0fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1281,4 +1281,219 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch) model="gpt-4", api_key=None, timezone_offset_minutes=480, - ) \ No newline at end of file + ) + + +# ===================== /v2/user/info tests ===================== + + +@pytest.mark.asyncio +async def test_user_info_v2_returns_user_profile(mocker): + """ + Test that /v2/user/info returns only the user profile (no keys or teams). + """ + from fastapi import Request + + from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_info_v2, + ) + + mock_prisma_client = mocker.MagicMock() + + mock_user = LiteLLM_UserTable( + user_id="test-user-123", + user_email="test@example.com", + user_role="internal_user", + teams=["team-1", "team-2"], + ) + + async def mock_get_data(*args, **kwargs): + if kwargs.get("user_id") == "test-user-123": + return mock_user + return None + + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_request = mocker.MagicMock(spec=Request) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="test-user-123", user_role="internal_user" + ) + + response = await user_info_v2( + user_id="test-user-123", + user_api_key_dict=mock_user_api_key_dict, + request=mock_request, + ) + + assert response.user_id == "test-user-123" + assert response.user_info.user_email == "test@example.com" + assert response.user_info.teams == ["team-1", "team-2"] + # Verify no keys or teams fields on the response + assert not hasattr(response, "keys") + assert not hasattr(response, "teams") + + +@pytest.mark.asyncio +async def test_user_info_v2_falls_back_to_caller_user_id(mocker): + """ + Test that /v2/user/info uses the caller's user_id when no user_id param is provided. + """ + from fastapi import Request + + from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_info_v2, + ) + + mock_prisma_client = mocker.MagicMock() + + mock_user = LiteLLM_UserTable( + user_id="caller-user-id", + user_email="caller@example.com", + ) + + async def mock_get_data(*args, **kwargs): + if kwargs.get("user_id") == "caller-user-id": + return mock_user + return None + + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_request = mocker.MagicMock(spec=Request) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="caller-user-id", user_role="internal_user" + ) + + response = await user_info_v2( + user_id=None, + user_api_key_dict=mock_user_api_key_dict, + request=mock_request, + ) + + assert response.user_id == "caller-user-id" + assert response.user_info.user_email == "caller@example.com" + + +@pytest.mark.asyncio +async def test_user_info_v2_404_for_unknown_user(mocker): + """ + Test that /v2/user/info returns 404 for a non-existent user. + """ + from fastapi import Request + + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_info_v2, + ) + + mock_prisma_client = mocker.MagicMock() + + async def mock_get_data(*args, **kwargs): + return None + + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_request = mocker.MagicMock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="admin-user", user_role="proxy_admin" + ) + + with pytest.raises(ProxyException) as exc_info: + await user_info_v2( + user_id="nonexistent-user", + user_api_key_dict=mock_user_api_key_dict, + request=mock_request, + ) + + assert exc_info.value.code == "404" + assert "nonexistent-user" in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_user_info_v2_400_when_no_user_id(mocker): + """ + Test that /v2/user/info returns 400 when no user_id is provided and + the caller's token has no user_id either. + """ + from fastapi import Request + + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_info_v2, + ) + + mock_prisma_client = mocker.MagicMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_request = mocker.MagicMock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_id=None, user_role="proxy_admin" + ) + + with pytest.raises(ProxyException) as exc_info: + await user_info_v2( + user_id=None, + user_api_key_dict=mock_user_api_key_dict, + request=mock_request, + ) + + assert exc_info.value.code == "400" + + +@pytest.mark.asyncio +async def test_user_info_v1_has_deprecation_header(mocker): + """ + Test that the old /user/info endpoint returns Deprecation headers. + """ + from fastapi import Request + + from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import user_info + + mock_prisma_client = mocker.MagicMock() + + mock_user = LiteLLM_UserTable( + user_id="test-user", + user_email="test@example.com", + teams=[], + ) + + async def mock_get_data(*args, **kwargs): + if kwargs.get("table_name") == "key": + return [] + elif kwargs.get("table_name") == "team": + return [] + elif kwargs.get("user_id") is not None: + return mock_user + return None + + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) + + mock_list_team = mocker.AsyncMock(return_value=None) + mocker.patch( + "litellm.proxy.management_endpoints.team_endpoints.list_team", + mock_list_team, + ) + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_request = mocker.MagicMock(spec=Request) + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="test-user", user_role="proxy_admin" + ) + + response = await user_info( + user_id="test-user", + user_api_key_dict=mock_user_api_key_dict, + request=mock_request, + ) + + # The response should now be a JSONResponse with deprecation headers + assert response.headers.get("Deprecation") == "true" + assert "successor-version" in response.headers.get("Link", "") \ No newline at end of file