User Headers X LiteLLM Users Mapping feature (#14485)

* Draft commit.

* user header mapping feature with backward compatibility with user_header_name field.

* user header mapping feature with backward compatibility with user_header_name field optimizations.

* Added unit tests.
This commit is contained in:
Boopesh Shanmugam
2025-09-12 11:49:37 -07:00
committed by GitHub
parent f8036a25a2
commit 8b338a4d8c
5 changed files with 206 additions and 6 deletions
+19
View File
@@ -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):
+36 -6
View File
@@ -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:
+39
View File
@@ -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:
+52
View File
@@ -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
@@ -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