Spaces:
Sleeping
Sleeping
""" | |
Handles Authentication Errors | |
""" | |
import asyncio | |
from typing import TYPE_CHECKING, Any, Optional, Union | |
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.proxy.db.exception_handler import PrismaDBExceptionHandler | |
from litellm.types.services import ServiceTypes | |
if TYPE_CHECKING: | |
from opentelemetry.trace import Span as _Span | |
Span = Union[_Span, Any] | |
else: | |
Span = Any | |
class UserAPIKeyAuthExceptionHandler: | |
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 ( | |
PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() | |
and PrismaDBExceptionHandler.is_database_connection_error(e) | |
): | |
# 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, | |
request_route=route, | |
) | |
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, | |
request_route=route, | |
) | |
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, | |
) | |