phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix check if user passed custom header

Browse source
This commit is contained in:
Ishaan Jaff 2024-07-16 16:36:18 -07:00
parent c0fd53a493
commit ceed3f655d
1 changed files with 29 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

33
litellm/proxy/auth/user_api_key_auth.py
View file

@ -135,13 +135,14 @@ async def user_api_key_auth(
header_key: str = headers.get("litellm_user_api_key", "") header_key: str = headers.get("litellm_user_api_key", "")
if request.headers.get(key=header_key) is not None: if request.headers.get(key=header_key) is not None:
api_key = request.headers.get(key=header_key) api_key = request.headers.get(key=header_key)
# if user wants to pass LiteLLM_Master_Key as a custom header, example pass litellm keys as X-LiteLLM-Key: Bearer sk-1234
custom_litellm_key_header_name = general_settings.get("litellm_key_header_name") custom_litellm_key_header_name = general_settings.get("litellm_key_header_name")
if custom_litellm_key_header_name is not None: if custom_litellm_key_header_name is not None:
# use this as the virtual key passed to litellm proxy api_key = get_api_key_from_custom_header(
passed_api_key = ( request=request,
request.headers.get(key=custom_litellm_key_header_name) or "" # type: ignore custom_litellm_key_header_name=custom_litellm_key_header_name,
) )
api_key = _get_bearer_token(api_key=passed_api_key)
parent_otel_span: Optional[Span] = None parent_otel_span: Optional[Span] = None
if open_telemetry_logger is not None: if open_telemetry_logger is not None:
@ -1272,3 +1273,27 @@ def _check_valid_ip(allowed_ips: Optional[List[str]], request: Request) -> bool:
return False return False
return True return True
def get_api_key_from_custom_header(
request: Request, custom_litellm_key_header_name: str
):
# use this as the virtual key passed to litellm proxy
custom_litellm_key_header_name = custom_litellm_key_header_name.lower()
verbose_proxy_logger.debug(
"searching for custom_litellm_key_header_name= %s",
custom_litellm_key_header_name,
)
custom_api_key = request.headers.get(custom_litellm_key_header_name)
if custom_api_key:
api_key = _get_bearer_token(api_key=custom_api_key)
verbose_proxy_logger.debug(
"Found custom API key using header: {}, setting api_key={}".format(
custom_litellm_key_header_name, api_key
)
)
else:
raise ValueError(
f"No LiteLLM Virtual Key pass. Please set header={custom_litellm_key_header_name}: Bearer <api_key>"
)
return api_key