tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-11 16:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix: Respect user_header_name property for budget selection and user identification (#11419)

Browse Source 
* Refactor get_end_user_id_from_request_body to support user ID retrieval from custom headers and multiple request body formats. Enhance tests to cover various scenarios including header precedence and fallback mechanisms.

* Refactor get_end_user_id_from_request_body function to accept request_body as the first parameter, improving clarity and flexibility. Update tests for compatibility and add new cases to ensure correct functionality across various request body formats.

* Update _user_api_key_auth_builder and user_api_key_auth to pass request object to get_end_user_id_from_request_body, enhancing user ID retrieval from request data.

* refactor(auth_utils.py): update get_end_user_id_from_request_body to accept request_headers instead of request, and adjust related function calls in user_api_key_auth and tests

* refactor(tests): update mock request handling in LLM pass-through endpoint tests

- Replaced the Request object with a Mock for better flexibility in testing.
- Enhanced mock setup to include user API key handling and virtual key retrieval.
- Updated test calls to reflect changes in mock request structure and added necessary patches for new dependencies.

* refactor(vertex_and_google_ai_studio_gemini.py): remove redundant variable declaration for url_context_metadata, linting error
This commit is contained in:
Cole McIntosh
2025-06-06 14:21:02 -07:00
committed by GitHub
parent f99e450d38
commit e191e72746
4 changed files with 238 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/auth/auth_utils.py
+36 -10
View File
@@ -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
litellm/proxy/auth/user_api_key_auth.py
+2 -2
View File
@@ -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
tests/local_testing/test_auth_utils.py
+157 -1
View File
@@ -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",
[
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
+43 -26
View File
@@ -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,