diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 66bd597755..46d6ac1f56 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1617,6 +1617,20 @@ class ConfigList(LiteLLMPydanticObjectBase): ) +class UserHeaderMapping(LiteLLMPydanticObjectBase): + """ + Map an incoming HTTP header to a LiteLLM user role. + """ + header_name: str + litellm_user_role: Literal[ + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.CUSTOMER, + ] + + model_config = { + "extra": "forbid", + } + class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ Documents all the fields supported by `general_settings` in config.yaml @@ -1721,6 +1735,11 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): default=None, description="Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through", ) + user_header_name: Optional[str] = Field( + None, + description="[DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings.", + ) + user_header_mappings: Optional[List[UserHeaderMapping]] = None class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 3261d15889..55f3f95539 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -473,6 +473,22 @@ def _has_user_setup_sso(): return sso_setup +def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: + """Return the header_name mapped to CUSTOMER role, if any (dict-based).""" + if not user_id_mapping: + return None + items = user_id_mapping if isinstance(user_id_mapping, list) else [user_id_mapping] + for item in items: + if not isinstance(item, dict): + continue + role = item.get("litellm_user_role") + header_name = item.get("header_name") + if role is None or not header_name: + continue + if str(role).lower() == str(LitellmUserRoles.CUSTOMER).lower(): + return header_name + return None + def get_end_user_id_from_request_body( request_body: dict, request_headers: Optional[dict] = None @@ -481,20 +497,34 @@ def get_end_user_id_from_request_body( # 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) + # Check 1 : Follow the user header mappings feature, if not found, then check for deprecated 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: Optional[str] = None - custom_header_name_to_check = general_settings.get(user_id_header_config_key) + # Prefer user mappings (new behavior) + user_id_mapping = general_settings.get("user_header_mappings", None) + if user_id_mapping: + custom_header_name_to_check = get_customer_user_header_from_mapping( + user_id_mapping + ) - if custom_header_name_to_check and isinstance(custom_header_name_to_check, str): + # Fallback to deprecated user_header_name if mapping did not specify + if not custom_header_name_to_check: + user_id_header_config_key = "user_header_name" + value = general_settings.get(user_id_header_config_key) + if isinstance(value, str) and value.strip() != "": + custom_header_name_to_check = value + + # If we have a header name to check, try to read it from request headers + if isinstance(custom_header_name_to_check, str): for header_name, header_value in request_headers.items(): if header_name.lower() == custom_header_name_to_check.lower(): user_id_from_header = header_value - if user_id_from_header.strip(): - return str(user_id_from_header) + user_id_str = str(user_id_from_header) if user_id_from_header is not None else "" + if user_id_str.strip(): + return user_id_str # Check 2: 'user' field in request_body (commonly OpenAI) if "user" in request_body and request_body["user"] is not None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 95f1eccffe..e82c8b2295 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -17,6 +17,7 @@ from litellm.proxy._types import ( SpecialHeaders, TeamCallbackMetadata, UserAPIKeyAuth, + LitellmUserRoles, ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.router import Router @@ -335,6 +336,22 @@ class LiteLLMProxyRequestSetup: return value return None + @staticmethod + def add_internal_user_from_user_mapping(general_settings: Optional[Dict], user_api_key_dict: UserAPIKeyAuth, headers: dict) -> UserAPIKeyAuth: + if general_settings is None: + return user_api_key_dict + user_header_mapping = general_settings.get("user_header_mappings") + if not user_header_mapping: + return user_api_key_dict + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(user_header_mapping) + if not header_name: + return user_api_key_dict + header_value = LiteLLMProxyRequestSetup._get_case_insensitive_header(headers, header_name) + if header_value: + user_api_key_dict.user_id = header_value + return user_api_key_dict + return user_api_key_dict + @staticmethod def get_user_from_headers( headers: dict, general_settings: Optional[Dict] = None @@ -428,6 +445,26 @@ class LiteLLMProxyRequestSetup: data["headers"] = _headers return data + @staticmethod + def get_internal_user_header_from_mapping(user_header_mapping) -> Optional[str]: + if not user_header_mapping: + return None + items = ( + user_header_mapping + if isinstance(user_header_mapping, list) + else [user_header_mapping] + ) + for item in items: + if not isinstance(item, dict): + continue + role = item.get("litellm_user_role") + header_name = item.get("header_name") + if role is None or not header_name: + continue + if str(role).lower() == str(LitellmUserRoles.INTERNAL_USER).lower(): + return header_name + return None + @staticmethod def add_litellm_data_for_backend_llm_call( *, @@ -726,6 +763,8 @@ async def add_litellm_data_to_request( # noqa: PLR0915 data=data, headers=_headers, user_api_key_dict=user_api_key_dict ) + user_api_key_dict = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(general_settings, user_api_key_dict, _headers) + # Parse user info from headers user = LiteLLMProxyRequestSetup.get_user_from_headers(_headers, general_settings) if user is not None: diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index b696eb2199..11261592c3 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -259,3 +259,55 @@ def test_get_model_from_request(request_data, expected_model): model = get_model_from_request(request_data, "/v1/files") assert model == ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"] + +def test_get_customer_user_header_from_mapping_returns_customer_header(): + from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping + + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, + ] + result = get_customer_user_header_from_mapping(mappings) + assert result == "X-OpenWebUI-User-Email" + + +def test_get_customer_user_header_from_mapping_no_customer_returns_none(): + from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping + + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} + ] + result = get_customer_user_header_from_mapping(mappings) + assert result is None + + # Also support a single mapping dict + single_mapping = {"header_name": "X-Only-Internal", "litellm_user_role": "internal_user"} + result = get_customer_user_header_from_mapping(single_mapping) + assert result is None + + +def test_get_internal_user_header_from_mapping_returns_internal_header(): + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, + ] + + result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) + assert result == "X-OpenWebUI-User-Id" + + +def test_get_internal_user_header_from_mapping_no_internal_returns_none(): + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + mappings = [ + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} + ] + result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) + assert result is None + + # Also support single mapping dict + single_mapping = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} + result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single_mapping) + assert result is None diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 817f19d8d7..bb95bd4109 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1059,3 +1059,63 @@ async def test_add_litellm_metadata_from_request_headers(): assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), "spend_logs_metadata should be the same as the headers" + +def test_get_internal_user_header_from_mapping_returns_expected_header(): + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, + ] + + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) + assert header_name == "X-OpenWebUI-User-Id" + + +def test_get_internal_user_header_from_mapping_none_when_absent(): + mappings = [ + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} + ] + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) + assert header_name is None + + single = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single) + assert header_name is None + + +def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + headers = {"X-OpenWebUI-User-Id": "internal-user-123"} + general_settings = { + "user_header_mappings": [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, + ] + } + + result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + general_settings, user_api_key_dict, headers + ) + + assert result is user_api_key_dict + assert user_api_key_dict.user_id == "internal-user-123" + + +def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchanged(): + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + None, user_api_key_dict, {"X-OpenWebUI-User-Id": "abc"} + ) + assert result is user_api_key_dict + assert user_api_key_dict.user_id is None + + general_settings = { + "user_header_mappings": [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} + ] + } + result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( + general_settings, user_api_key_dict, {"Other": "value"} + ) + assert result is user_api_key_dict + assert user_api_key_dict.user_id is None