mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-14 04:25:53 +00:00
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:
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user