diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a26f43d9dd..f170664053 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -499,17 +499,43 @@ def _has_user_setup_sso(): return sso_setup -def get_end_user_id_from_request_body(request_body: dict) -> Optional[str]: - # openai - check 'user' +def get_end_user_id_from_request_body(request_body: dict, request_headers: Optional[dict] = None) -> Optional[str]: + # Import general_settings here to avoid potential circular import issues at module level + # and to ensure it's fetched at runtime. + from litellm.proxy.proxy_server import general_settings + + # Check 1: Custom Header from general_settings.user_header_name (only if request_headers is provided) + # User query: "system not respecting user_header_name property" + # This implies the key in general_settings is 'user_header_name'. + if request_headers is not None: + user_id_header_config_key = "user_header_name" + + custom_header_name_to_check = general_settings.get(user_id_header_config_key) + + if custom_header_name_to_check and isinstance(custom_header_name_to_check, str): + user_id_from_header = request_headers.get(custom_header_name_to_check) + if user_id_from_header is not None and user_id_from_header.strip(): + return str(user_id_from_header) + + # Check 2: 'user' field in request_body (commonly OpenAI) if "user" in request_body and request_body["user"] is not None: - return str(request_body["user"]) - # anthropic - check 'litellm_metadata' - end_user_id = request_body.get("litellm_metadata", {}).get("user", None) - if end_user_id: - return str(end_user_id) - metadata = request_body.get("metadata") - if metadata and "user_id" in metadata and metadata["user_id"] is not None: - return str(metadata["user_id"]) + user_from_body_user_field = request_body["user"] + return str(user_from_body_user_field) + + # Check 3: 'litellm_metadata.user' in request_body (commonly Anthropic) + litellm_metadata = request_body.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + user_from_litellm_metadata = litellm_metadata.get("user") + if user_from_litellm_metadata is not None: + return str(user_from_litellm_metadata) + + # Check 4: 'metadata.user_id' in request_body (another common pattern) + metadata_dict = request_body.get("metadata") + if isinstance(metadata_dict, dict): + user_id_from_metadata_field = metadata_dict.get("user_id") + if user_id_from_metadata_field is not None: + return str(user_id_from_metadata_field) + return None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 31f33d92e8..73d6bf35df 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -587,7 +587,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _end_user_object = None end_user_params = {} - end_user_id = get_end_user_id_from_request_body(request_data) + end_user_id = get_end_user_id_from_request_body(request_data, dict(request.headers)) if end_user_id: try: end_user_params["end_user_id"] = end_user_id @@ -1140,7 +1140,7 @@ async def user_api_key_auth( custom_litellm_key_header=custom_litellm_key_header, ) - end_user_id = get_end_user_id_from_request_body(request_data) + end_user_id = get_end_user_id_from_request_body(request_data, dict(request.headers)) if end_user_id is not None: user_api_key_auth_obj.end_user_id = end_user_id diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index 6d73195e38..e2ddbca603 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -72,12 +72,168 @@ def test_configurable_clientside_parameters( def test_get_end_user_id_from_request_body_always_returns_str(): from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + from fastapi import Request + from unittest.mock import MagicMock + # Create a mock Request object + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + request_body = {"user": 123} - end_user_id = get_end_user_id_from_request_body(request_body) + end_user_id = get_end_user_id_from_request_body(request_body, dict(mock_request.headers)) assert end_user_id == "123" assert isinstance(end_user_id, str) + +@pytest.mark.parametrize( + "headers, general_settings_config, request_body, expected_user_id", + [ + # Test 1: user_header_name configured and header present + ( + {"X-User-ID": "header-user-123"}, + {"user_header_name": "X-User-ID"}, + {"user": "body-user-456"}, + "header-user-123" # Header should take precedence + ), + # Test 2: user_header_name configured but header not present, fallback to body + ( + {}, + {"user_header_name": "X-User-ID"}, + {"user": "body-user-456"}, + "body-user-456" # Should fall back to body + ), + # Test 3: user_header_name not configured, should use body + ( + {"X-User-ID": "header-user-123"}, + {}, + {"user": "body-user-456"}, + "body-user-456" # Should ignore header when not configured + ), + # Test 4: user_header_name configured, header present, but no body user + ( + {"X-Custom-User": "header-only-user"}, + {"user_header_name": "X-Custom-User"}, + {"model": "gpt-4"}, + "header-only-user" # Should use header + ), + # Test 5: user_header_name configured but header is empty string + ( + {"X-User-ID": ""}, + {"user_header_name": "X-User-ID"}, + {"user": "body-user-456"}, + "body-user-456" # Should fall back to body when header is empty + ), + # Test 6: user_header_name configured with case-insensitive header + ( + {"x-user-id": "lowercase-header-user"}, + {"user_header_name": "x-user-id"}, + {"user": "body-user-456"}, + "lowercase-header-user" + ), + # Test 7: user_header_name configured but set to None + ( + {"X-User-ID": "header-user-123"}, + {"user_header_name": None}, + {"user": "body-user-456"}, + "body-user-456" # Should fall back to body when header name is None + ), + # Test 8: user_header_name is not a string + ( + {"X-User-ID": "header-user-123"}, + {"user_header_name": 123}, + {"user": "body-user-456"}, + "body-user-456" # Should fall back to body when header name is not a string + ), + # Test 9: Multiple fallback sources - litellm_metadata + ( + {}, + {"user_header_name": "X-User-ID"}, + {"litellm_metadata": {"user": "litellm-user-789"}}, + "litellm-user-789" + ), + # Test 10: Multiple fallback sources - metadata.user_id + ( + {}, + {"user_header_name": "X-User-ID"}, + {"metadata": {"user_id": "metadata-user-999"}}, + "metadata-user-999" + ), + # Test 11: Header takes precedence over all body sources + ( + {"X-User-ID": "header-priority"}, + {"user_header_name": "X-User-ID"}, + { + "user": "body-user", + "litellm_metadata": {"user": "litellm-user"}, + "metadata": {"user_id": "metadata-user"} + }, + "header-priority" + ), + ] +) +def test_get_end_user_id_from_request_body_with_user_header_name( + headers, general_settings_config, request_body, expected_user_id +): + """Test that get_end_user_id_from_request_body respects user_header_name property""" + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + from fastapi import Request + from unittest.mock import MagicMock, patch + + # Create a mock Request object with headers + mock_request = MagicMock(spec=Request) + mock_request.headers = headers + + # Mock general_settings at the proxy_server module level + with patch('litellm.proxy.proxy_server.general_settings', general_settings_config): + end_user_id = get_end_user_id_from_request_body(request_body, dict(mock_request.headers)) + assert end_user_id == expected_user_id + + +def test_get_end_user_id_from_request_body_no_user_found(): + """Test that function returns None when no user ID is found anywhere""" + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + from fastapi import Request + from unittest.mock import MagicMock, patch + + # Create a mock Request object with no relevant headers + mock_request = MagicMock(spec=Request) + mock_request.headers = {"X-Other-Header": "some-value"} + + # Mock general_settings with user_header_name that doesn't match headers + general_settings_config = {"user_header_name": "X-User-ID"} + + # Request body with no user identifiers + request_body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]} + + with patch('litellm.proxy.proxy_server.general_settings', general_settings_config): + end_user_id = get_end_user_id_from_request_body(request_body, dict(mock_request.headers)) + assert end_user_id is None + + +def test_get_end_user_id_from_request_body_backwards_compatibility(): + """Test that function works with just request_body parameter (backwards compatibility)""" + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + + # Test with just request_body - should work like before + request_body = {"user": "test-user-123"} + end_user_id = get_end_user_id_from_request_body(request_body) + assert end_user_id == "test-user-123" + + # Test with litellm_metadata + request_body = {"litellm_metadata": {"user": "litellm-user-456"}} + end_user_id = get_end_user_id_from_request_body(request_body) + assert end_user_id == "litellm-user-456" + + # Test with metadata.user_id + request_body = {"metadata": {"user_id": "metadata-user-789"}} + end_user_id = get_end_user_id_from_request_body(request_body) + assert end_user_id == "metadata-user-789" + + # Test with no user - should return None + request_body = {"model": "gpt-4"} + end_user_id = get_end_user_id_from_request_body(request_body) + assert end_user_id is None + @pytest.mark.parametrize( "request_data, expected_model", [ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 9dcde87a59..cf88dfb14f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -220,17 +220,14 @@ class TestVertexAIPassThroughHandler: endpoint = f"/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/gemini-1.5-flash:generateContent" # Mock request - mock_request = Request( - scope={ - "type": "http", - "method": "POST", - "path": endpoint, - "headers": [ - (b"Authorization", b"Bearer test-creds"), - (b"Content-Type", b"application/json"), - ], - } - ) + mock_request = Mock() + mock_request.method = "POST" + mock_request.headers = { + "Authorization": "Bearer test-creds", + "Content-Type": "application/json", + } + mock_request.url = Mock() + mock_request.url.path = endpoint # Mock response mock_response = Response() @@ -246,16 +243,28 @@ class TestVertexAIPassThroughHandler: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.vertex_llm_base._get_token_and_url" ) as mock_get_token, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_create_route: + ) as mock_create_route, mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key" + ) as mock_get_virtual_key, mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" + ) as mock_user_auth: + # Setup mocks mock_ensure_token.return_value = ("test-auth-header", test_project) mock_get_token.return_value = (test_token, "") + mock_get_virtual_key.return_value = "Bearer test-key" + mock_user_auth.return_value = {"api_key": "test-key"} + + # Mock create_pass_through_route to return a function that returns a mock response + mock_endpoint_func = AsyncMock(return_value={"status": "success"}) + mock_create_route.return_value = mock_endpoint_func # Call the route try: - await vertex_proxy_route( + result = await vertex_proxy_route( endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict={"api_key": "test-key"}, ) except Exception as e: print(f"Error: {e}") @@ -492,17 +501,14 @@ class TestVertexAIDiscoveryPassThroughHandler: endpoint = f"/v1/projects/{vertex_project}/locations/{vertex_location}/dataStores/default/servingConfigs/default:search" # Mock request - mock_request = Request( - scope={ - "type": "http", - "method": "POST", - "path": endpoint, - "headers": [ - (b"Authorization", b"Bearer test-creds"), - (b"Content-Type", b"application/json"), - ], - } - ) + mock_request = Mock() + mock_request.method = "POST" + mock_request.headers = { + "Authorization": "Bearer test-creds", + "Content-Type": "application/json", + } + mock_request.url = Mock() + mock_request.url.path = endpoint # Mock response mock_response = Response() @@ -518,13 +524,24 @@ class TestVertexAIDiscoveryPassThroughHandler: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.vertex_llm_base._get_token_and_url" ) as mock_get_token, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_create_route: + ) as mock_create_route, mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_virtual_key" + ) as mock_get_virtual_key, mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" + ) as mock_user_auth: + # Setup mocks mock_ensure_token.return_value = ("test-auth-header", test_project) mock_get_token.return_value = (test_token, "") + mock_get_virtual_key.return_value = "Bearer test-key" + mock_user_auth.return_value = {"api_key": "test-key"} + + # Mock create_pass_through_route to return a function that returns a mock response + mock_endpoint_func = AsyncMock(return_value={"status": "success"}) + mock_create_route.return_value = mock_endpoint_func # Call the route try: - await vertex_discovery_proxy_route( + result = await vertex_discovery_proxy_route( endpoint=endpoint, request=mock_request, fastapi_response=mock_response,