mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
UserAPIKeyAuthExceptionHandler
This commit is contained in:
Ishaan Jaff
parent
ff61ce6751
commit
4c87084ff7
3 changed files with 136 additions and 91 deletions
|
|
@ -1009,8 +1009,6 @@ async def get_key_object(
|
|||
)
|
||||
|
||||
return _response
|
||||
except DB_CONNECTION_ERROR_TYPES as e:
|
||||
return await _handle_failed_db_connection_for_get_key_object(e=e)
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
raise Exception(
|
||||
|
|
@ -1018,46 +1016,6 @@ async def get_key_object(
|
|||
)
|
||||
|
||||
|
||||
async def _handle_failed_db_connection_for_get_key_object(
|
||||
e: Exception,
|
||||
) -> UserAPIKeyAuth:
|
||||
"""
|
||||
Handles httpx.ConnectError when reading a Virtual Key from LiteLLM DB
|
||||
|
||||
Use this if you don't want failed DB queries to block LLM API reqiests
|
||||
|
||||
Returns:
|
||||
- UserAPIKeyAuth: If general_settings.allow_requests_on_db_unavailable is True
|
||||
|
||||
Raises:
|
||||
- Orignal Exception in all other cases
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
litellm_proxy_admin_name,
|
||||
proxy_logging_obj,
|
||||
)
|
||||
|
||||
# If this flag is on, requests failing to connect to the DB will be allowed
|
||||
if general_settings.get("allow_requests_on_db_unavailable", False) is True:
|
||||
# log this as a DB failure on prometheus
|
||||
proxy_logging_obj.service_logging_obj.service_failure_hook(
|
||||
service=ServiceTypes.DB,
|
||||
call_type="get_key_object",
|
||||
error=e,
|
||||
duration=0.0,
|
||||
)
|
||||
|
||||
return UserAPIKeyAuth(
|
||||
key_name="failed-to-connect-to-db",
|
||||
token="failed-to-connect-to-db",
|
||||
user_id=litellm_proxy_admin_name,
|
||||
)
|
||||
else:
|
||||
# raise the original exception, the wrapper on `get_key_object` handles logging db failure to prometheus
|
||||
raise e
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_org_object(
|
||||
org_id: str,
|
||||
|
|
|
|||
130
litellm/proxy/auth/auth_exception_handler.py
Normal file
130
litellm/proxy/auth/auth_exception_handler.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
"""
|
||||
Handles Authentication Errors
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import _get_request_ip_address
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
Span = _Span
|
||||
else:
|
||||
Span = Any
|
||||
|
||||
|
||||
class UserAPIKeyAuthExceptionHandler:
|
||||
|
||||
@staticmethod
|
||||
async def _handle_authentication_error(
|
||||
e: Exception,
|
||||
request: Request,
|
||||
request_data: dict,
|
||||
route: str,
|
||||
parent_otel_span: Optional[Span],
|
||||
api_key: str,
|
||||
) -> UserAPIKeyAuth:
|
||||
"""
|
||||
Handles Connection Errors when reading a Virtual Key from LiteLLM DB
|
||||
Use this if you don't want failed DB queries to block LLM API reqiests
|
||||
|
||||
Reliability scenarios this covers:
|
||||
- DB is down and having an outage
|
||||
- Unable to read / recover a key from the DB
|
||||
|
||||
Returns:
|
||||
- UserAPIKeyAuth: If general_settings.allow_requests_on_db_unavailable is True
|
||||
|
||||
Raises:
|
||||
- Orignal Exception in all other cases
|
||||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
litellm_proxy_admin_name,
|
||||
proxy_logging_obj,
|
||||
)
|
||||
|
||||
if UserAPIKeyAuthExceptionHandler.should_allow_request_on_db_unavailable():
|
||||
# log this as a DB failure on prometheus
|
||||
proxy_logging_obj.service_logging_obj.service_failure_hook(
|
||||
service=ServiceTypes.DB,
|
||||
call_type="get_key_object",
|
||||
error=e,
|
||||
duration=0.0,
|
||||
)
|
||||
|
||||
return UserAPIKeyAuth(
|
||||
key_name="failed-to-connect-to-db",
|
||||
token="failed-to-connect-to-db",
|
||||
user_id=litellm_proxy_admin_name,
|
||||
)
|
||||
else:
|
||||
# raise the exception to the caller
|
||||
requester_ip = _get_request_ip_address(
|
||||
request=request,
|
||||
use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False),
|
||||
)
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format(
|
||||
str(e),
|
||||
requester_ip,
|
||||
),
|
||||
extra={"requester_ip": requester_ip},
|
||||
)
|
||||
|
||||
# Log this exception to OTEL, Datadog etc
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
parent_otel_span=parent_otel_span,
|
||||
api_key=api_key,
|
||||
)
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
route=route,
|
||||
)
|
||||
)
|
||||
|
||||
if isinstance(e, litellm.BudgetExceededError):
|
||||
raise ProxyException(
|
||||
message=e.message,
|
||||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param=None,
|
||||
code=400,
|
||||
)
|
||||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "detail", f"Authentication Error({str(e)})"),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED),
|
||||
)
|
||||
elif isinstance(e, ProxyException):
|
||||
raise e
|
||||
raise ProxyException(
|
||||
message="Authentication Error, " + str(e),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def should_allow_request_on_db_unavailable() -> bool:
|
||||
"""
|
||||
Returns True if the request should be allowed to proceed despite the DB connection error
|
||||
"""
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
if general_settings.get("allow_requests_on_db_unavailable", False) is True:
|
||||
return True
|
||||
return False
|
||||
|
|
@ -26,7 +26,6 @@ from litellm.proxy._types import *
|
|||
from litellm.proxy.auth.auth_checks import (
|
||||
_cache_key_object,
|
||||
_get_user_role,
|
||||
_handle_failed_db_connection_for_get_key_object,
|
||||
_is_user_proxy_admin,
|
||||
_virtual_key_max_budget_check,
|
||||
_virtual_key_soft_budget_check,
|
||||
|
|
@ -38,6 +37,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
get_user_object,
|
||||
is_valid_fallback_model,
|
||||
)
|
||||
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_get_request_ip_address,
|
||||
get_end_user_id_from_request_body,
|
||||
|
|
@ -675,9 +675,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
if (
|
||||
prisma_client is None
|
||||
): # if both master key + user key submitted, and user key != master key, and no db connected, raise an error
|
||||
return await _handle_failed_db_connection_for_get_key_object(
|
||||
e=Exception("No connected db.")
|
||||
)
|
||||
raise Exception("No connected db.")
|
||||
|
||||
## check for cache hit (In-Memory Cache)
|
||||
_user_role = None
|
||||
|
|
@ -1018,55 +1016,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
|||
else:
|
||||
raise Exception()
|
||||
except Exception as e:
|
||||
requester_ip = _get_request_ip_address(
|
||||
return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
|
||||
e=e,
|
||||
request=request,
|
||||
use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False),
|
||||
)
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format(
|
||||
str(e),
|
||||
requester_ip,
|
||||
),
|
||||
extra={"requester_ip": requester_ip},
|
||||
)
|
||||
|
||||
# Log this exception to OTEL, Datadog etc
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
parent_otel_span=parent_otel_span,
|
||||
api_key=api_key,
|
||||
)
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
route=route,
|
||||
)
|
||||
)
|
||||
|
||||
if isinstance(e, litellm.BudgetExceededError):
|
||||
raise ProxyException(
|
||||
message=e.message,
|
||||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param=None,
|
||||
code=400,
|
||||
)
|
||||
if isinstance(e, HTTPException):
|
||||
raise ProxyException(
|
||||
message=getattr(e, "detail", f"Authentication Error({str(e)})"),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED),
|
||||
)
|
||||
elif isinstance(e, ProxyException):
|
||||
raise e
|
||||
raise ProxyException(
|
||||
message="Authentication Error, " + str(e),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param=getattr(e, "param", "None"),
|
||||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
)
|
||||
|
||||
|
||||
@tracer.wrap()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue