From 4c87084ff7bd99966a33b6337ef58a5a2f11d9b0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 14:07:14 -0700 Subject: [PATCH 01/28] UserAPIKeyAuthExceptionHandler --- litellm/proxy/auth/auth_checks.py | 42 ------ litellm/proxy/auth/auth_exception_handler.py | 130 +++++++++++++++++++ litellm/proxy/auth/user_api_key_auth.py | 55 +------- 3 files changed, 136 insertions(+), 91 deletions(-) create mode 100644 litellm/proxy/auth/auth_exception_handler.py diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 80cfb03de4..98685e1a7c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py new file mode 100644 index 0000000000..88a94fd5b9 --- /dev/null +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b78619ae65..a2850ca294 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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() From 59040167ac7dd3e1864f76a51a8b29dff6ef96b0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 14:40:11 -0700 Subject: [PATCH 02/28] fix ProxyErrorTypes --- litellm/proxy/_types.py | 53 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 49 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 220a0d5ddb..54612868cc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2067,16 +2067,64 @@ class SpendCalculateRequest(LiteLLMPydanticObjectBase): class ProxyErrorTypes(str, enum.Enum): budget_exceeded = "budget_exceeded" + """ + Object was over budget + """ + + token_not_found_in_db = "token_not_found_in_db" + """ + Requested token was not found in the database + """ + key_model_access_denied = "key_model_access_denied" + """ + Key does not have access to the model + """ + team_model_access_denied = "team_model_access_denied" + """ + Team does not have access to the model + """ + user_model_access_denied = "user_model_access_denied" + """ + User does not have access to the model + """ + expired_key = "expired_key" + """ + Key has expired + """ + auth_error = "auth_error" + """ + General authentication error + """ + internal_server_error = "internal_server_error" + """ + Internal server error + """ + bad_request_error = "bad_request_error" + """ + Bad request error + """ + not_found_error = "not_found_error" - validation_error = "bad_request_error" + """ + Not found error + """ + + validation_error = "validation_error" + """ + Validation error + """ + cache_ping_error = "cache_ping_error" + """ + Cache ping error + """ @classmethod def get_model_access_error_type_for_object( @@ -2093,9 +2141,6 @@ class ProxyErrorTypes(str, enum.Enum): return cls.user_model_access_denied -DB_CONNECTION_ERROR_TYPES = (httpx.ConnectError, httpx.ReadError, httpx.ReadTimeout) - - class SSOUserDefinedValues(TypedDict): models: List[str] user_id: str From ce49e2721789cfaa4a236103cef90742f7e51bfb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 15:44:13 -0700 Subject: [PATCH 03/28] fixes for auth checks --- litellm/proxy/_types.py | 7 +++ litellm/proxy/auth/auth_checks.py | 47 +++++++++--------- litellm/proxy/auth/auth_exception_handler.py | 23 ++++++++- litellm/proxy/auth/user_api_key_auth.py | 50 ++++++++------------ 4 files changed, 71 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 54612868cc..c0fb0750eb 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2141,6 +2141,13 @@ class ProxyErrorTypes(str, enum.Enum): return cls.user_model_access_denied +DB_CONNECTION_ERROR_TYPES = ( + httpx.ConnectError, + httpx.ReadError, + httpx.ReadTimeout, +) + + class SSOUserDefinedValues(TypedDict): models: List[str] user_id: str diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 98685e1a7c..9e71147706 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -987,33 +987,34 @@ async def get_key_object( ) # else, check db - try: - _valid_token: Optional[BaseModel] = await prisma_client.get_data( - token=hashed_token, - table_name="combined_view", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + _valid_token: Optional[BaseModel] = await prisma_client.get_data( + token=hashed_token, + table_name="combined_view", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + if _valid_token is None: + raise ProxyException( + message="Key doesn't exist in db. key={}. Create key via `/key/generate` call.".format( + hashed_token + ), + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, ) - if _valid_token is None: - raise Exception + _response = UserAPIKeyAuth(**_valid_token.model_dump(exclude_none=True)) - _response = UserAPIKeyAuth(**_valid_token.model_dump(exclude_none=True)) + # save the key object to cache + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=_response, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) - # save the key object to cache - await _cache_key_object( - hashed_token=hashed_token, - user_api_key_obj=_response, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - return _response - except Exception: - traceback.print_exc() - raise Exception( - f"Key doesn't exist in db. key={hashed_token}. Create key via `/key/generate` call." - ) + return _response @log_db_metrics diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 88a94fd5b9..5b314ef931 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -9,7 +9,12 @@ 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._types import ( + DB_CONNECTION_ERROR_TYPES, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, +) from litellm.proxy.auth.auth_utils import _get_request_ip_address from litellm.types.services import ServiceTypes @@ -52,7 +57,10 @@ class UserAPIKeyAuthExceptionHandler: proxy_logging_obj, ) - if UserAPIKeyAuthExceptionHandler.should_allow_request_on_db_unavailable(): + if ( + UserAPIKeyAuthExceptionHandler.should_allow_request_on_db_unavailable() + and UserAPIKeyAuthExceptionHandler.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, @@ -128,3 +136,14 @@ class UserAPIKeyAuthExceptionHandler: if general_settings.get("allow_requests_on_db_unavailable", False) is True: return True return False + + @staticmethod + def is_database_connection_error(e: Exception) -> bool: + """ + Returns True if the exception is from a database outage / connection error + """ + import prisma + + return isinstance(e, DB_CONNECTION_ERROR_TYPES) or isinstance( + e, prisma.errors.PrismaError + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a2850ca294..06fa560866 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -683,37 +683,25 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 api_key = hash_token(token=api_key) if valid_token is None: - try: - valid_token = await get_key_object( - hashed_token=api_key, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - # update end-user params on valid token - # These can change per request - it's important to update them here - valid_token.end_user_id = end_user_params.get("end_user_id") - valid_token.end_user_tpm_limit = end_user_params.get( - "end_user_tpm_limit" - ) - valid_token.end_user_rpm_limit = end_user_params.get( - "end_user_rpm_limit" - ) - valid_token.allowed_model_region = end_user_params.get( - "allowed_model_region" - ) - # update key budget with temp budget increase - valid_token = _update_key_budget_with_temp_budget_increase( - valid_token - ) # updating it here, allows all downstream reporting / checks to use the updated budget - except Exception: - verbose_logger.info( - "litellm.proxy.auth.user_api_key_auth.py::user_api_key_auth() - Unable to find token={} in cache or `LiteLLM_VerificationTokenTable`. Defaulting 'valid_token' to None'".format( - api_key - ) - ) - valid_token = None + valid_token = await get_key_object( + hashed_token=api_key, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + # update end-user params on valid token + # These can change per request - it's important to update them here + valid_token.end_user_id = end_user_params.get("end_user_id") + valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") + valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") + valid_token.allowed_model_region = end_user_params.get( + "allowed_model_region" + ) + # update key budget with temp budget increase + valid_token = _update_key_budget_with_temp_budget_increase( + valid_token + ) # updating it here, allows all downstream reporting / checks to use the updated budget if valid_token is None: raise Exception( From c6d5793bf6077ac94b7884ee1fde97521ac6ab52 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 17:50:27 -0700 Subject: [PATCH 04/28] add toxi proxy tests to ci/cd --- .circleci/config.yml | 63 ++++++++++++ .../setup_toxi_proxy.py | 96 +++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 tests/proxy_reliability_tests/setup_toxi_proxy.py diff --git a/.circleci/config.yml b/.circleci/config.yml index b93a9d81e8..87432c7a24 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -470,6 +470,61 @@ jobs: paths: - litellm_proxy_security_tests_coverage.xml - litellm_proxy_security_tests_coverage + litellm_proxy_reliability_tests: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + steps: + - checkout + - run: + name: Show git commit hash + command: | + echo "Git commit hash: $CIRCLE_SHA1" + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "requests==2.31.0" + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-asyncio==0.21.1" + pip install "pytest-cov==5.0.0" + - run: + name: Setup Toxiproxy + command: | + python tests/proxy_reliability_tests/setup_toxiproxy.py + - run: + name: Run prisma ./docker/entrypoint.sh + command: | + set +e + chmod +x docker/entrypoint.sh + ./docker/entrypoint.sh + set -e + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest tests/proxy_security_tests --cov=litellm --cov-report=xml -vv -x -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml litellm_proxy_security_tests_coverage.xml + mv .coverage litellm_proxy_security_tests_coverage + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - litellm_proxy_security_tests_coverage.xml + - litellm_proxy_security_tests_coverage litellm_proxy_unit_testing: # Runs all tests with the "proxy", "key", "jwt" filenames docker: - image: cimg/python:3.11 @@ -2451,6 +2506,12 @@ workflows: only: - main - /litellm_.*/ + - litellm_proxy_reliability_tests: + filters: + branches: + only: + - main + - /litellm_.*/ - litellm_assistants_api_testing: filters: branches: @@ -2592,6 +2653,7 @@ workflows: - caching_unit_tests - litellm_proxy_unit_testing - litellm_proxy_security_tests + - litellm_proxy_reliability_tests - langfuse_logging_unit_tests - local_testing - litellm_assistants_api_testing @@ -2657,6 +2719,7 @@ workflows: - e2e_ui_testing - litellm_proxy_unit_testing - litellm_proxy_security_tests + - litellm_proxy_reliability_tests - installing_litellm_on_python - installing_litellm_on_python_3_13 - proxy_logging_guardrails_model_info_tests diff --git a/tests/proxy_reliability_tests/setup_toxi_proxy.py b/tests/proxy_reliability_tests/setup_toxi_proxy.py new file mode 100644 index 0000000000..3625b0bd3a --- /dev/null +++ b/tests/proxy_reliability_tests/setup_toxi_proxy.py @@ -0,0 +1,96 @@ +import os +import platform +import subprocess +import time +import requests +import stat + + +def download_toxiproxy(): + # Determine system architecture and OS + system = platform.system().lower() + machine = platform.machine().lower() + + # Map architecture names + arch_map = { + "x86_64": "amd64", + "amd64": "amd64", + "arm64": "arm64", + "aarch64": "arm64", + } + + arch = arch_map.get(machine, machine) + + # Construct download URL (using latest version 2.5.0) + base_url = "https://github.com/Shopify/toxiproxy/releases/download/v2.5.0/" + if system == "linux": + filename = f"toxiproxy-server-linux-{arch}" + cli_filename = f"toxiproxy-cli-linux-{arch}" + elif system == "darwin": + filename = f"toxiproxy-server-darwin-{arch}" + cli_filename = f"toxiproxy-cli-darwin-{arch}" + else: + raise Exception("Unsupported operating system") + + # Download server + response = requests.get(f"{base_url}{filename}") + with open("toxiproxy-server", "wb") as f: + f.write(response.content) + + # Download CLI + response = requests.get(f"{base_url}{cli_filename}") + with open("toxiproxy-cli", "wb") as f: + f.write(response.content) + + # Make files executable + os.chmod("toxiproxy-server", stat.S_IRWXU) + os.chmod("toxiproxy-cli", stat.S_IRWXU) + + +def setup_toxiproxy(): + # Start toxiproxy-server + server_process = subprocess.Popen(["./toxiproxy-server"]) + + # Wait for server to start + time.sleep(2) + + # Create proxy + subprocess.run( + [ + "./toxiproxy-cli", + "create", + "postgres_proxy", + "-l", + "127.0.0.1:6666", + "-u", + "ep-dry-paper-a69g2y1q-pooler.us-west-2.aws.neon.tech:5432", + ] + ) + + return server_process + + +def main(): + try: + # Download ToxiProxy binaries + download_toxiproxy() + + # Setup ToxiProxy + server_process = setup_toxiproxy() + + print("ToxiProxy setup completed successfully!") + print("Proxy 'postgres_proxy' created and listening on 127.0.0.1:6666") + + # Keep the script running to maintain the proxy + try: + server_process.wait() + except KeyboardInterrupt: + server_process.terminate() + print("\nToxiProxy server stopped") + + except Exception as e: + print(f"Error: {e}") + + +if __name__ == "__main__": + main() From 7b09d88680430d2fe5fd79c0f57e6079fa08113b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 17:52:12 -0700 Subject: [PATCH 05/28] fix setup --- tests/proxy_reliability_tests/setup_toxi_proxy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/proxy_reliability_tests/setup_toxi_proxy.py b/tests/proxy_reliability_tests/setup_toxi_proxy.py index 3625b0bd3a..1338445f3a 100644 --- a/tests/proxy_reliability_tests/setup_toxi_proxy.py +++ b/tests/proxy_reliability_tests/setup_toxi_proxy.py @@ -59,11 +59,11 @@ def setup_toxiproxy(): [ "./toxiproxy-cli", "create", - "postgres_proxy", "-l", "127.0.0.1:6666", "-u", "ep-dry-paper-a69g2y1q-pooler.us-west-2.aws.neon.tech:5432", + "postgres_proxy", ] ) From 34c3825d1306c095730bdf1dd52f1e9c25c2b917 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 17:53:30 -0700 Subject: [PATCH 06/28] fix path --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 87432c7a24..ea887bf9de 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -496,7 +496,7 @@ jobs: - run: name: Setup Toxiproxy command: | - python tests/proxy_reliability_tests/setup_toxiproxy.py + python tests/proxy_reliability_tests/setup_toxi_proxy.py - run: name: Run prisma ./docker/entrypoint.sh command: | From 53a586e87626c8816ad8febde009e53be5890c09 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 17:59:38 -0700 Subject: [PATCH 07/28] TOXI_PROXY_DATABASE_URL --- .circleci/config.yml | 75 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 58 insertions(+), 17 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index ea887bf9de..15eacf9f51 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -471,47 +471,88 @@ jobs: - litellm_proxy_security_tests_coverage.xml - litellm_proxy_security_tests_coverage litellm_proxy_reliability_tests: - docker: - - image: cimg/python:3.11 - auth: - username: ${DOCKERHUB_USERNAME} - password: ${DOCKERHUB_PASSWORD} + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge working_directory: ~/project steps: - checkout - run: - name: Show git commit hash + name: Install Docker CLI (In case it's not already installed) command: | - echo "Git commit hash: $CIRCLE_SHA1" + sudo apt-get update + sudo apt-get install -y docker-ce docker-ce-cli containerd.io + - run: + name: Install Python 3.9 + command: | + curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh + bash miniconda.sh -b -p $HOME/miniconda + export PATH="$HOME/miniconda/bin:$PATH" + conda init bash + source ~/.bashrc + conda create -n myenv python=3.9 -y + conda activate myenv + python --version - run: name: Install Dependencies command: | + pip install "pytest==7.3.1" + pip install "pytest-asyncio==0.21.1" + pip install aiohttp python -m pip install --upgrade pip python -m pip install -r requirements.txt - pip install "requests==2.31.0" pip install "pytest==7.3.1" pip install "pytest-retry==1.6.3" + pip install "pytest-mock==3.12.0" pip install "pytest-asyncio==0.21.1" - pip install "pytest-cov==5.0.0" + pip install "assemblyai==0.37.0" - run: name: Setup Toxiproxy command: | python tests/proxy_reliability_tests/setup_toxi_proxy.py - run: - name: Run prisma ./docker/entrypoint.sh + name: Build Docker image + command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + - run: + name: Run Docker container + # intentionally give bad redis credentials here + # the OTEL test - should get this as a trace command: | - set +e - chmod +x docker/entrypoint.sh - ./docker/entrypoint.sh - set -e - # Run pytest and generate JUnit XML report + docker run -d \ + -p 4000:4000 \ + -e DATABASE_URL=$TOXI_PROXY_DATABASE_URL \ + -e STORE_MODEL_IN_DB="True" \ + -e LITELLM_MASTER_KEY="sk-1234" \ + -e LITELLM_LICENSE=$LITELLM_LICENSE \ + --name my-app \ + -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ + my-app:latest \ + --config /app/config.yaml \ + --port 4000 \ + --detailed_debug \ + - run: + name: Install curl and dockerize + command: | + sudo apt-get update + sudo apt-get install -y curl + sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz + sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz + sudo rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Start outputting logs + command: docker logs -f my-app + background: true + - run: + name: Wait for app to be ready + command: dockerize -wait http://localhost:4000 -timeout 5m - run: name: Run tests command: | pwd ls - python -m pytest tests/proxy_security_tests --cov=litellm --cov-report=xml -vv -x -v --junitxml=test-results/junit.xml --durations=5 - no_output_timeout: 120m + python -m pytest -vv tests/store_model_in_db_tests -x --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: + 120m - run: name: Rename the coverage files command: | From bf7241abd1cb6d92899df8796fbdb149df01af72 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 18:02:01 -0700 Subject: [PATCH 08/28] litellm_proxy_reliability_tests --- .circleci/config.yml | 193 ++++++++++++++++++++++--------------------- 1 file changed, 97 insertions(+), 96 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 15eacf9f51..61f98405a7 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -470,102 +470,6 @@ jobs: paths: - litellm_proxy_security_tests_coverage.xml - litellm_proxy_security_tests_coverage - litellm_proxy_reliability_tests: - machine: - image: ubuntu-2204:2023.10.1 - resource_class: xlarge - working_directory: ~/project - steps: - - checkout - - run: - name: Install Docker CLI (In case it's not already installed) - command: | - sudo apt-get update - sudo apt-get install -y docker-ce docker-ce-cli containerd.io - - run: - name: Install Python 3.9 - command: | - curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh - bash miniconda.sh -b -p $HOME/miniconda - export PATH="$HOME/miniconda/bin:$PATH" - conda init bash - source ~/.bashrc - conda create -n myenv python=3.9 -y - conda activate myenv - python --version - - run: - name: Install Dependencies - command: | - pip install "pytest==7.3.1" - pip install "pytest-asyncio==0.21.1" - pip install aiohttp - python -m pip install --upgrade pip - python -m pip install -r requirements.txt - pip install "pytest==7.3.1" - pip install "pytest-retry==1.6.3" - pip install "pytest-mock==3.12.0" - pip install "pytest-asyncio==0.21.1" - pip install "assemblyai==0.37.0" - - run: - name: Setup Toxiproxy - command: | - python tests/proxy_reliability_tests/setup_toxi_proxy.py - - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . - - run: - name: Run Docker container - # intentionally give bad redis credentials here - # the OTEL test - should get this as a trace - command: | - docker run -d \ - -p 4000:4000 \ - -e DATABASE_URL=$TOXI_PROXY_DATABASE_URL \ - -e STORE_MODEL_IN_DB="True" \ - -e LITELLM_MASTER_KEY="sk-1234" \ - -e LITELLM_LICENSE=$LITELLM_LICENSE \ - --name my-app \ - -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ - my-app:latest \ - --config /app/config.yaml \ - --port 4000 \ - --detailed_debug \ - - run: - name: Install curl and dockerize - command: | - sudo apt-get update - sudo apt-get install -y curl - sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz - sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz - sudo rm dockerize-linux-amd64-v0.6.1.tar.gz - - run: - name: Start outputting logs - command: docker logs -f my-app - background: true - - run: - name: Wait for app to be ready - command: dockerize -wait http://localhost:4000 -timeout 5m - - run: - name: Run tests - command: | - pwd - ls - python -m pytest -vv tests/store_model_in_db_tests -x --junitxml=test-results/junit.xml --durations=5 - no_output_timeout: - 120m - - run: - name: Rename the coverage files - command: | - mv coverage.xml litellm_proxy_security_tests_coverage.xml - mv .coverage litellm_proxy_security_tests_coverage - # Store test results - - store_test_results: - path: test-results - - persist_to_workspace: - root: . - paths: - - litellm_proxy_security_tests_coverage.xml - - litellm_proxy_security_tests_coverage litellm_proxy_unit_testing: # Runs all tests with the "proxy", "key", "jwt" filenames docker: - image: cimg/python:3.11 @@ -1984,6 +1888,103 @@ jobs: no_output_timeout: 120m # Clean up first container + litellm_proxy_reliability_tests: + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge + working_directory: ~/project + steps: + - checkout + - run: + name: Install Docker CLI (In case it's not already installed) + command: | + sudo apt-get update + sudo apt-get install -y docker-ce docker-ce-cli containerd.io + - run: + name: Install Python 3.9 + command: | + curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh + bash miniconda.sh -b -p $HOME/miniconda + export PATH="$HOME/miniconda/bin:$PATH" + conda init bash + source ~/.bashrc + conda create -n myenv python=3.9 -y + conda activate myenv + python --version + - run: + name: Install Dependencies + command: | + pip install "pytest==7.3.1" + pip install "pytest-asyncio==0.21.1" + pip install aiohttp + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-mock==3.12.0" + pip install "pytest-asyncio==0.21.1" + pip install "assemblyai==0.37.0" + - run: + name: Setup Toxiproxy + command: | + python tests/proxy_reliability_tests/setup_toxi_proxy.py + - run: + name: Build Docker image + command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + - run: + name: Run Docker container + # intentionally give bad redis credentials here + # the OTEL test - should get this as a trace + command: | + docker run -d \ + -p 4000:4000 \ + -e DATABASE_URL=$TOXI_PROXY_DATABASE_URL \ + -e STORE_MODEL_IN_DB="True" \ + -e LITELLM_MASTER_KEY="sk-1234" \ + -e LITELLM_LICENSE=$LITELLM_LICENSE \ + --name my-app \ + -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ + my-app:latest \ + --config /app/config.yaml \ + --port 4000 \ + --detailed_debug \ + - run: + name: Install curl and dockerize + command: | + sudo apt-get update + sudo apt-get install -y curl + sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz + sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz + sudo rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Start outputting logs + command: docker logs -f my-app + background: true + - run: + name: Wait for app to be ready + command: dockerize -wait http://localhost:4000 -timeout 5m + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/store_model_in_db_tests -x --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: + 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml litellm_proxy_security_tests_coverage.xml + mv .coverage litellm_proxy_security_tests_coverage + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - litellm_proxy_security_tests_coverage.xml + - litellm_proxy_security_tests_coverage + proxy_build_from_pip_tests: # Change from docker to machine executor From 83b41f95e769a1f19d33d1b8ca934311153e6dd5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 18:05:41 -0700 Subject: [PATCH 09/28] Setup Toxiproxy --- .circleci/config.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 61f98405a7..b7dc08bced 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1927,7 +1927,8 @@ jobs: - run: name: Setup Toxiproxy command: | - python tests/proxy_reliability_tests/setup_toxi_proxy.py + python tests/proxy_reliability_tests/setup_toxi_proxy.py & + sleep 5 - run: name: Build Docker image command: docker build -t my-app:latest -f ./docker/Dockerfile.database . From 6f138c79a744cb1f9ea7f44089968f132e737580 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 18:19:11 -0700 Subject: [PATCH 10/28] run toxi proxy tests --- .circleci/config.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index b7dc08bced..2577862036 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1934,8 +1934,6 @@ jobs: command: docker build -t my-app:latest -f ./docker/Dockerfile.database . - run: name: Run Docker container - # intentionally give bad redis credentials here - # the OTEL test - should get this as a trace command: | docker run -d \ -p 4000:4000 \ From 9e2d230339c001b894956a841689bf115176927e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 18:23:52 -0700 Subject: [PATCH 11/28] litellm_proxy_reliability_tests --- .circleci/config.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 2577862036..8e376b6dc2 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2552,7 +2552,8 @@ workflows: branches: only: - main - - /litellm_.*/ + - /litellm_stability_.*/ + - litellm_stable_release_branch - litellm_assistants_api_testing: filters: branches: From 53d9e33e782da6051418eff380e6b507fcafa329 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 18:59:26 -0700 Subject: [PATCH 12/28] fix setup toxi proxy --- .circleci/config.yml | 1 + tests/proxy_reliability_tests/setup_toxi_proxy.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 8e376b6dc2..04799ddd75 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1942,6 +1942,7 @@ jobs: -e LITELLM_MASTER_KEY="sk-1234" \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ --name my-app \ + --network=host \ -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ my-app:latest \ --config /app/config.yaml \ diff --git a/tests/proxy_reliability_tests/setup_toxi_proxy.py b/tests/proxy_reliability_tests/setup_toxi_proxy.py index 1338445f3a..b44298aca5 100644 --- a/tests/proxy_reliability_tests/setup_toxi_proxy.py +++ b/tests/proxy_reliability_tests/setup_toxi_proxy.py @@ -60,7 +60,7 @@ def setup_toxiproxy(): "./toxiproxy-cli", "create", "-l", - "127.0.0.1:6666", + "0.0.0.0:6666", "-u", "ep-dry-paper-a69g2y1q-pooler.us-west-2.aws.neon.tech:5432", "postgres_proxy", From 438655858210efc9cbc9d3d594bbe66494c92c66 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 19:11:13 -0700 Subject: [PATCH 13/28] litellm_proxy_reliability_tests --- .circleci/config.yml | 137 ++++++++++++------------------------------- 1 file changed, 39 insertions(+), 98 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 04799ddd75..e4d6af8f0e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -686,6 +686,44 @@ jobs: paths: - llm_translation_coverage.xml - llm_translation_coverage + litellm_proxy_reliability_tests: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + + steps: + - checkout + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-cov==5.0.0" + pip install "pytest-asyncio==0.21.1" + pip install "respx==0.21.1" + # Run pytest and generate JUnit XML report + - run: + name: Setup Toxiproxy + command: | + python tests/proxy_reliability_tests/setup_toxi_proxy.py & + sleep 5 + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/proxy_reliability_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml litellm_proxy_reliability_tests_coverage.xml + mv .coverage litellm_proxy_reliability_tests_coverage mcp_testing: docker: - image: cimg/python:3.11 @@ -1888,104 +1926,7 @@ jobs: no_output_timeout: 120m # Clean up first container - litellm_proxy_reliability_tests: - machine: - image: ubuntu-2204:2023.10.1 - resource_class: xlarge - working_directory: ~/project - steps: - - checkout - - run: - name: Install Docker CLI (In case it's not already installed) - command: | - sudo apt-get update - sudo apt-get install -y docker-ce docker-ce-cli containerd.io - - run: - name: Install Python 3.9 - command: | - curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh - bash miniconda.sh -b -p $HOME/miniconda - export PATH="$HOME/miniconda/bin:$PATH" - conda init bash - source ~/.bashrc - conda create -n myenv python=3.9 -y - conda activate myenv - python --version - - run: - name: Install Dependencies - command: | - pip install "pytest==7.3.1" - pip install "pytest-asyncio==0.21.1" - pip install aiohttp - python -m pip install --upgrade pip - python -m pip install -r requirements.txt - pip install "pytest==7.3.1" - pip install "pytest-retry==1.6.3" - pip install "pytest-mock==3.12.0" - pip install "pytest-asyncio==0.21.1" - pip install "assemblyai==0.37.0" - - run: - name: Setup Toxiproxy - command: | - python tests/proxy_reliability_tests/setup_toxi_proxy.py & - sleep 5 - - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . - - run: - name: Run Docker container - command: | - docker run -d \ - -p 4000:4000 \ - -e DATABASE_URL=$TOXI_PROXY_DATABASE_URL \ - -e STORE_MODEL_IN_DB="True" \ - -e LITELLM_MASTER_KEY="sk-1234" \ - -e LITELLM_LICENSE=$LITELLM_LICENSE \ - --name my-app \ - --network=host \ - -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ - my-app:latest \ - --config /app/config.yaml \ - --port 4000 \ - --detailed_debug \ - - run: - name: Install curl and dockerize - command: | - sudo apt-get update - sudo apt-get install -y curl - sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz - sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz - sudo rm dockerize-linux-amd64-v0.6.1.tar.gz - - run: - name: Start outputting logs - command: docker logs -f my-app - background: true - - run: - name: Wait for app to be ready - command: dockerize -wait http://localhost:4000 -timeout 5m - - run: - name: Run tests - command: | - pwd - ls - python -m pytest -vv tests/store_model_in_db_tests -x --junitxml=test-results/junit.xml --durations=5 - no_output_timeout: - 120m - - run: - name: Rename the coverage files - command: | - mv coverage.xml litellm_proxy_security_tests_coverage.xml - mv .coverage litellm_proxy_security_tests_coverage - # Store test results - - store_test_results: - path: test-results - - persist_to_workspace: - root: . - paths: - - litellm_proxy_security_tests_coverage.xml - - litellm_proxy_security_tests_coverage - - + proxy_build_from_pip_tests: # Change from docker to machine executor machine: From 9d10befa09f5cd16420195abaef99b9795bb9bc1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 19:16:34 -0700 Subject: [PATCH 14/28] test_litellm_proxy_server_config_no_general_settings --- .../test_using_litellm_db_outage.py | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 tests/proxy_reliability_tests/test_using_litellm_db_outage.py diff --git a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py new file mode 100644 index 0000000000..745a14bdc6 --- /dev/null +++ b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py @@ -0,0 +1,74 @@ +import asyncio +import os +import subprocess +import sys +import time +import traceback + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import litellm + +import os +import subprocess +import time + +import pytest +import requests + + +def test_litellm_proxy_server_config_no_general_settings(): + # Install the litellm[proxy] package + # Start the server + try: + subprocess.run(["pip", "install", "litellm[proxy]"]) + subprocess.run(["pip", "install", "litellm[extra_proxy]"]) + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + server_process = subprocess.Popen( + [ + "python", + "-m", + "litellm.proxy.proxy_cli", + "--config", + config_fp, + ] + ) + + # Allow some time for the server to start + time.sleep(60) # Adjust the sleep time if necessary + + # Send a request to the /health/liveliness endpoint + response = requests.get("http://localhost:4000/health/liveliness") + + # Check if the response is successful + assert response.status_code == 200 + assert response.json() == "I'm alive!" + + # Test /chat/completions + response = requests.post( + "http://localhost:4000/chat/completions", + headers={"Authorization": "Bearer 1234567890"}, + json={ + "model": "test_openai_models", + "messages": [{"role": "user", "content": "Hello, how are you?"}], + }, + ) + + assert response.status_code == 200 + + except ImportError: + pytest.fail("Failed to import litellm.proxy_server") + except requests.ConnectionError: + pytest.fail("Failed to connect to the server") + finally: + # Shut down the server + server_process.terminate() + server_process.wait() + + # Additional assertions can be added here + assert True From b4e745323a49e630b0b1f98bdefcce62690e2d04 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 19:21:51 -0700 Subject: [PATCH 15/28] add test config --- litellm/proxy/proxy_config.yaml | 40 +++---------------- .../test_configs/test_config.yaml | 9 +++++ .../test_using_litellm_db_outage.py | 2 +- 3 files changed, 16 insertions(+), 35 deletions(-) create mode 100644 tests/proxy_reliability_tests/test_configs/test_config.yaml diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 26ce6cb8f8..4912a35f89 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,37 +1,9 @@ model_list: - - model_name: gpt-3.5-turbo-end-user-test + - model_name: fake-openai-endpoint litellm_params: - model: azure/chatgpt-v-2 - api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ - api_version: "2023-05-15" - api_key: os.environ/AZURE_API_KEY + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ - - -mcp_tools: - - name: "get_current_time" - description: "Get the current time" - input_schema: { - "type": "object", - "properties": { - "format": { - "type": "string", - "description": "The format of the time to return", - "enum": ["short"] - } - } - } - handler: "mcp_tools.get_current_time" - - name: "get_current_date" - description: "Get the current date" - input_schema: { - "type": "object", - "properties": { - "format": { - "type": "string", - "description": "The format of the date to return", - "enum": ["short"] - } - } - } - handler: "mcp_tools.get_current_date" +general_settings: + allow_requests_on_db_unavailable: True \ No newline at end of file diff --git a/tests/proxy_reliability_tests/test_configs/test_config.yaml b/tests/proxy_reliability_tests/test_configs/test_config.yaml new file mode 100644 index 0000000000..4912a35f89 --- /dev/null +++ b/tests/proxy_reliability_tests/test_configs/test_config.yaml @@ -0,0 +1,9 @@ +model_list: + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + +general_settings: + allow_requests_on_db_unavailable: True \ No newline at end of file diff --git a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py index 745a14bdc6..acaaed05f3 100644 --- a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py +++ b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py @@ -28,7 +28,7 @@ def test_litellm_proxy_server_config_no_general_settings(): subprocess.run(["pip", "install", "litellm[proxy]"]) subprocess.run(["pip", "install", "litellm[extra_proxy]"]) filepath = os.path.dirname(os.path.abspath(__file__)) - config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" + config_fp = f"{filepath}/test_configs/test_config.yaml" server_process = subprocess.Popen( [ "python", From 6572ba7a0ef9ddb24345dde179610365d237e571 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 19:25:47 -0700 Subject: [PATCH 16/28] fix startup --- .../proxy_reliability_tests/test_using_litellm_db_outage.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py index acaaed05f3..52d3ec6222 100644 --- a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py +++ b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py @@ -4,7 +4,6 @@ import subprocess import sys import time import traceback - import pytest sys.path.insert( @@ -29,6 +28,10 @@ def test_litellm_proxy_server_config_no_general_settings(): subprocess.run(["pip", "install", "litellm[extra_proxy]"]) filepath = os.path.dirname(os.path.abspath(__file__)) config_fp = f"{filepath}/test_configs/test_config.yaml" + + # Set DATABASE_URL environment variable + os.environ["DATABASE_URL"] = os.getenv("TOXI_PROXY_DATABASE_URL") + server_process = subprocess.Popen( [ "python", From 0a401ee468de6a13c905c8b41c66a73b95d5b1c5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 19:27:15 -0700 Subject: [PATCH 17/28] test_litellm_proxy_server_config_no_general_settings --- .../test_using_litellm_db_outage.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py index 52d3ec6222..ea9c7e3f3a 100644 --- a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py +++ b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py @@ -19,18 +19,20 @@ import time import pytest import requests +TEST_MASTER_KEY = "sk-1234" + def test_litellm_proxy_server_config_no_general_settings(): # Install the litellm[proxy] package # Start the server try: subprocess.run(["pip", "install", "litellm[proxy]"]) - subprocess.run(["pip", "install", "litellm[extra_proxy]"]) filepath = os.path.dirname(os.path.abspath(__file__)) config_fp = f"{filepath}/test_configs/test_config.yaml" # Set DATABASE_URL environment variable os.environ["DATABASE_URL"] = os.getenv("TOXI_PROXY_DATABASE_URL") + os.environ["LITELLM_MASTER_KEY"] = TEST_MASTER_KEY server_process = subprocess.Popen( [ @@ -72,6 +74,3 @@ def test_litellm_proxy_server_config_no_general_settings(): # Shut down the server server_process.terminate() server_process.wait() - - # Additional assertions can be added here - assert True From 6e5d2b1ac77ef48fb4f65ef131799c8ca183bcf4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Mar 2025 23:14:44 -0700 Subject: [PATCH 18/28] handle failed db connections --- tests/proxy_unit_tests/test_auth_checks.py | 33 ---------------------- 1 file changed, 33 deletions(-) diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index 0eb1a38755..7695306c87 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -13,9 +13,6 @@ sys.path.insert( ) # Adds the parent directory to the system path import pytest, litellm import httpx -from litellm.proxy.auth.auth_checks import ( - _handle_failed_db_connection_for_get_key_object, -) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import get_end_user_object from litellm.caching.caching import DualCache @@ -78,36 +75,6 @@ async def test_get_end_user_object(customer_spend, customer_budget): ) -@pytest.mark.asyncio -async def test_handle_failed_db_connection(): - """ - Test cases: - 1. When allow_requests_on_db_unavailable=True -> return UserAPIKeyAuth - 2. When allow_requests_on_db_unavailable=False -> raise original error - """ - from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name - - # Test case 1: allow_requests_on_db_unavailable=True - general_settings["allow_requests_on_db_unavailable"] = True - mock_error = httpx.ConnectError("Failed to connect to DB") - - result = await _handle_failed_db_connection_for_get_key_object(e=mock_error) - - assert isinstance(result, UserAPIKeyAuth) - assert result.key_name == "failed-to-connect-to-db" - assert result.token == "failed-to-connect-to-db" - assert result.user_id == litellm_proxy_admin_name - - # Test case 2: allow_requests_on_db_unavailable=False - general_settings["allow_requests_on_db_unavailable"] = False - - with pytest.raises(httpx.ConnectError) as exc_info: - await _handle_failed_db_connection_for_get_key_object(e=mock_error) - print("_handle_failed_db_connection_for_get_key_object got exception", exc_info) - - assert str(exc_info.value) == "Failed to connect to DB" - - @pytest.mark.parametrize( "model, expect_to_work", [ From 1812ce4a54e5c15c1206667b812c7fde924ed0d2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 14:43:03 -0700 Subject: [PATCH 19/28] undo config.yml changes --- .circleci/config.yml | 49 +------------------------------------------- 1 file changed, 1 insertion(+), 48 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 43de7e69e8..5a058144f6 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -710,44 +710,6 @@ jobs: paths: - llm_translation_coverage.xml - llm_translation_coverage - litellm_proxy_reliability_tests: - docker: - - image: cimg/python:3.11 - auth: - username: ${DOCKERHUB_USERNAME} - password: ${DOCKERHUB_PASSWORD} - working_directory: ~/project - - steps: - - checkout - - run: - name: Install Dependencies - command: | - python -m pip install --upgrade pip - python -m pip install -r requirements.txt - pip install "pytest==7.3.1" - pip install "pytest-retry==1.6.3" - pip install "pytest-cov==5.0.0" - pip install "pytest-asyncio==0.21.1" - pip install "respx==0.21.1" - # Run pytest and generate JUnit XML report - - run: - name: Setup Toxiproxy - command: | - python tests/proxy_reliability_tests/setup_toxi_proxy.py & - sleep 5 - - run: - name: Run tests - command: | - pwd - ls - python -m pytest -vv tests/proxy_reliability_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 - no_output_timeout: 120m - - run: - name: Rename the coverage files - command: | - mv coverage.xml litellm_proxy_reliability_tests_coverage.xml - mv .coverage litellm_proxy_reliability_tests_coverage mcp_testing: docker: - image: cimg/python:3.11 @@ -1968,7 +1930,7 @@ jobs: no_output_timeout: 120m # Clean up first container - + proxy_build_from_pip_tests: # Change from docker to machine executor machine: @@ -2536,13 +2498,6 @@ workflows: only: - main - /litellm_.*/ - - litellm_proxy_reliability_tests: - filters: - branches: - only: - - main - - /litellm_stability_.*/ - - litellm_stable_release_branch - litellm_assistants_api_testing: filters: branches: @@ -2684,7 +2639,6 @@ workflows: - caching_unit_tests - litellm_proxy_unit_testing - litellm_proxy_security_tests - - litellm_proxy_reliability_tests - langfuse_logging_unit_tests - local_testing - litellm_assistants_api_testing @@ -2750,7 +2704,6 @@ workflows: - e2e_ui_testing - litellm_proxy_unit_testing - litellm_proxy_security_tests - - litellm_proxy_reliability_tests - installing_litellm_on_python - installing_litellm_on_python_3_13 - proxy_logging_guardrails_model_info_tests From 4948673e3572b16f95057e881d219b598a6a9d04 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 14:51:33 -0700 Subject: [PATCH 20/28] fix test changes --- .../setup_toxi_proxy.py | 96 ------------------- .../test_configs/test_config.yaml | 9 -- .../test_using_litellm_db_outage.py | 76 --------------- 3 files changed, 181 deletions(-) delete mode 100644 tests/proxy_reliability_tests/setup_toxi_proxy.py delete mode 100644 tests/proxy_reliability_tests/test_configs/test_config.yaml delete mode 100644 tests/proxy_reliability_tests/test_using_litellm_db_outage.py diff --git a/tests/proxy_reliability_tests/setup_toxi_proxy.py b/tests/proxy_reliability_tests/setup_toxi_proxy.py deleted file mode 100644 index b44298aca5..0000000000 --- a/tests/proxy_reliability_tests/setup_toxi_proxy.py +++ /dev/null @@ -1,96 +0,0 @@ -import os -import platform -import subprocess -import time -import requests -import stat - - -def download_toxiproxy(): - # Determine system architecture and OS - system = platform.system().lower() - machine = platform.machine().lower() - - # Map architecture names - arch_map = { - "x86_64": "amd64", - "amd64": "amd64", - "arm64": "arm64", - "aarch64": "arm64", - } - - arch = arch_map.get(machine, machine) - - # Construct download URL (using latest version 2.5.0) - base_url = "https://github.com/Shopify/toxiproxy/releases/download/v2.5.0/" - if system == "linux": - filename = f"toxiproxy-server-linux-{arch}" - cli_filename = f"toxiproxy-cli-linux-{arch}" - elif system == "darwin": - filename = f"toxiproxy-server-darwin-{arch}" - cli_filename = f"toxiproxy-cli-darwin-{arch}" - else: - raise Exception("Unsupported operating system") - - # Download server - response = requests.get(f"{base_url}{filename}") - with open("toxiproxy-server", "wb") as f: - f.write(response.content) - - # Download CLI - response = requests.get(f"{base_url}{cli_filename}") - with open("toxiproxy-cli", "wb") as f: - f.write(response.content) - - # Make files executable - os.chmod("toxiproxy-server", stat.S_IRWXU) - os.chmod("toxiproxy-cli", stat.S_IRWXU) - - -def setup_toxiproxy(): - # Start toxiproxy-server - server_process = subprocess.Popen(["./toxiproxy-server"]) - - # Wait for server to start - time.sleep(2) - - # Create proxy - subprocess.run( - [ - "./toxiproxy-cli", - "create", - "-l", - "0.0.0.0:6666", - "-u", - "ep-dry-paper-a69g2y1q-pooler.us-west-2.aws.neon.tech:5432", - "postgres_proxy", - ] - ) - - return server_process - - -def main(): - try: - # Download ToxiProxy binaries - download_toxiproxy() - - # Setup ToxiProxy - server_process = setup_toxiproxy() - - print("ToxiProxy setup completed successfully!") - print("Proxy 'postgres_proxy' created and listening on 127.0.0.1:6666") - - # Keep the script running to maintain the proxy - try: - server_process.wait() - except KeyboardInterrupt: - server_process.terminate() - print("\nToxiProxy server stopped") - - except Exception as e: - print(f"Error: {e}") - - -if __name__ == "__main__": - main() diff --git a/tests/proxy_reliability_tests/test_configs/test_config.yaml b/tests/proxy_reliability_tests/test_configs/test_config.yaml deleted file mode 100644 index 4912a35f89..0000000000 --- a/tests/proxy_reliability_tests/test_configs/test_config.yaml +++ /dev/null @@ -1,9 +0,0 @@ -model_list: - - model_name: fake-openai-endpoint - litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ - -general_settings: - allow_requests_on_db_unavailable: True \ No newline at end of file diff --git a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py b/tests/proxy_reliability_tests/test_using_litellm_db_outage.py deleted file mode 100644 index ea9c7e3f3a..0000000000 --- a/tests/proxy_reliability_tests/test_using_litellm_db_outage.py +++ /dev/null @@ -1,76 +0,0 @@ -import asyncio -import os -import subprocess -import sys -import time -import traceback -import pytest - -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path - -import litellm - -import os -import subprocess -import time - -import pytest -import requests - -TEST_MASTER_KEY = "sk-1234" - - -def test_litellm_proxy_server_config_no_general_settings(): - # Install the litellm[proxy] package - # Start the server - try: - subprocess.run(["pip", "install", "litellm[proxy]"]) - filepath = os.path.dirname(os.path.abspath(__file__)) - config_fp = f"{filepath}/test_configs/test_config.yaml" - - # Set DATABASE_URL environment variable - os.environ["DATABASE_URL"] = os.getenv("TOXI_PROXY_DATABASE_URL") - os.environ["LITELLM_MASTER_KEY"] = TEST_MASTER_KEY - - server_process = subprocess.Popen( - [ - "python", - "-m", - "litellm.proxy.proxy_cli", - "--config", - config_fp, - ] - ) - - # Allow some time for the server to start - time.sleep(60) # Adjust the sleep time if necessary - - # Send a request to the /health/liveliness endpoint - response = requests.get("http://localhost:4000/health/liveliness") - - # Check if the response is successful - assert response.status_code == 200 - assert response.json() == "I'm alive!" - - # Test /chat/completions - response = requests.post( - "http://localhost:4000/chat/completions", - headers={"Authorization": "Bearer 1234567890"}, - json={ - "model": "test_openai_models", - "messages": [{"role": "user", "content": "Hello, how are you?"}], - }, - ) - - assert response.status_code == 200 - - except ImportError: - pytest.fail("Failed to import litellm.proxy_server") - except requests.ConnectionError: - pytest.fail("Failed to connect to the server") - finally: - # Shut down the server - server_process.terminate() - server_process.wait() From f8caebc54b28dda058ec42ea638d344a5662a784 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 14:55:40 -0700 Subject: [PATCH 21/28] is_database_connection_error --- litellm/proxy/_types.py | 4 ++++ litellm/proxy/auth/auth_exception_handler.py | 10 +++++++--- litellm/proxy/auth/user_api_key_auth.py | 9 ++++++--- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c0fb0750eb..104d1e7338 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2070,6 +2070,10 @@ class ProxyErrorTypes(str, enum.Enum): """ Object was over budget """ + no_db_connection = "no_db_connection" + """ + No database connection + """ token_not_found_in_db = "token_not_found_in_db" """ diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 5b314ef931..c1a546b569 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -144,6 +144,10 @@ class UserAPIKeyAuthExceptionHandler: """ import prisma - return isinstance(e, DB_CONNECTION_ERROR_TYPES) or isinstance( - e, prisma.errors.PrismaError - ) + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return True + if isinstance(e, prisma.errors.PrismaError): + return True + if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection: + return True + return False diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 06fa560866..85ee76ff31 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -675,7 +675,12 @@ 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 - raise Exception("No connected db.") + raise ProxyException( + message="No connected db.", + type=ProxyErrorTypes.no_db_connection, + code=400, + param=None, + ) ## check for cache hit (In-Memory Cache) _user_role = None @@ -1001,8 +1006,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, start_time=start_time, ) - else: - raise Exception() except Exception as e: return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( e=e, From 8bd2081dec15c2255e295302687a1e1c4ce413e3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 15:41:40 -0700 Subject: [PATCH 22/28] fix get_key_object --- litellm/proxy/auth/auth_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9e71147706..dc0ec116f5 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -996,7 +996,7 @@ async def get_key_object( if _valid_token is None: raise ProxyException( - message="Key doesn't exist in db. key={}. Create key via `/key/generate` call.".format( + message="Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( hashed_token ), type=ProxyErrorTypes.token_not_found_in_db, From ff33ed020cc72fd4fa5541f58d4974f731b11006 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 15:45:58 -0700 Subject: [PATCH 23/28] fix auth checks --- litellm/proxy/auth/auth_checks.py | 3 --- litellm/proxy/auth/user_api_key_auth.py | 1 - 2 files changed, 4 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dc0ec116f5..7f75272bbd 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,7 +11,6 @@ Run checks for: import asyncio import re import time -import traceback from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, cast from fastapi import Request, status @@ -23,7 +22,6 @@ from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, RBAC_ROLES, CallInfo, LiteLLM_EndUserTable, @@ -45,7 +43,6 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router -from litellm.types.services import ServiceTypes from .auth_checks_organization import organization_role_based_access_check diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 85ee76ff31..b58353bf05 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -39,7 +39,6 @@ from litellm.proxy.auth.auth_checks import ( ) 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, get_request_route, is_pass_through_provider_route, From 23aa7f81b562b4cd96079e05733acb56b7c0bd1b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 15:53:33 -0700 Subject: [PATCH 24/28] fix ProxyException --- litellm/proxy/auth/auth_checks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 7f75272bbd..efbfe8d90c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -993,7 +993,7 @@ async def get_key_object( if _valid_token is None: raise ProxyException( - message="Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( + message="Authentication Error, Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( hashed_token ), type=ProxyErrorTypes.token_not_found_in_db, From 15b1a8afb0c85191c0cc2229c2b2cd8be71cabef Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 18:27:39 -0700 Subject: [PATCH 25/28] test_is_database_connection_error_prisma_errors --- .../proxy/auth/test_auth_exception_handler.py | 161 ++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 tests/litellm/proxy/auth/test_auth_exception_handler.py diff --git a/tests/litellm/proxy/auth/test_auth_exception_handler.py b/tests/litellm/proxy/auth/test_auth_exception_handler.py new file mode 100644 index 0000000000..ca0e8c88c4 --- /dev/null +++ b/tests/litellm/proxy/auth/test_auth_exception_handler.py @@ -0,0 +1,161 @@ +import asyncio +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest +from fastapi import HTTPException, Request, status +from prisma import errors as prisma_errors +from prisma.errors import ( + ClientNotConnectedError, + DataError, + ForeignKeyViolationError, + HTTPClientClosedError, + MissingRequiredValueError, + PrismaError, + RawQueryError, + RecordNotFoundError, + TableNotFoundError, + UniqueViolationError, +) + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler + + +# Test is_database_connection_error method +@pytest.mark.parametrize( + "prisma_error", + [ + PrismaError(), + DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + ForeignKeyViolationError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + MissingRequiredValueError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + TableNotFoundError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + RecordNotFoundError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + HTTPClientClosedError(), + ClientNotConnectedError(), + ], +) +def test_is_database_connection_error_prisma_errors(prisma_error): + """ + Test that all Prisma errors are considered database connection errors + """ + handler = UserAPIKeyAuthExceptionHandler() + assert handler.is_database_connection_error(prisma_error) == True + + +def test_is_database_connection_generic_errors(): + """ + Test non-Prisma error cases for database connection checking + """ + handler = UserAPIKeyAuthExceptionHandler() + + # Test with ProxyException (DB connection) + db_proxy_exception = ProxyException( + message="DB Connection Error", + type=ProxyErrorTypes.no_db_connection, + param="test-param", + ) + assert handler.is_database_connection_error(db_proxy_exception) == True + + # Test with non-DB error + regular_exception = Exception("Regular error") + assert handler.is_database_connection_error(regular_exception) == False + + +# Test should_allow_request_on_db_unavailable method +@patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, +) +def test_should_allow_request_on_db_unavailable_true(): + handler = UserAPIKeyAuthExceptionHandler() + assert handler.should_allow_request_on_db_unavailable() == True + + +@patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, +) +def test_should_allow_request_on_db_unavailable_false(): + handler = UserAPIKeyAuthExceptionHandler() + assert handler.should_allow_request_on_db_unavailable() == False + + +# Test _handle_authentication_error method +@pytest.mark.asyncio +async def test_handle_authentication_error_db_unavailable(): + handler = UserAPIKeyAuthExceptionHandler() + + # Mock request and other dependencies + mock_request = MagicMock() + mock_request_data = {} + mock_route = "/test" + mock_span = None + mock_api_key = "test-key" + + # Test with DB connection error when requests are allowed + with patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ): + db_error = prisma_errors.PrismaError() + result = await handler._handle_authentication_error( + db_error, + mock_request, + mock_request_data, + mock_route, + mock_span, + mock_api_key, + ) + assert result.key_name == "failed-to-connect-to-db" + assert result.token == "failed-to-connect-to-db" + + +@pytest.mark.asyncio +async def test_handle_authentication_error_budget_exceeded(): + handler = UserAPIKeyAuthExceptionHandler() + + # Mock request and other dependencies + mock_request = MagicMock() + mock_request_data = {} + mock_route = "/test" + mock_span = None + mock_api_key = "test-key" + + # Test with budget exceeded error + with pytest.raises(ProxyException) as exc_info: + from litellm.exceptions import BudgetExceededError + + budget_error = BudgetExceededError( + message="Budget exceeded", current_cost=100, max_budget=100 + ) + await handler._handle_authentication_error( + budget_error, + mock_request, + mock_request_data, + mock_route, + mock_span, + mock_api_key, + ) + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded From 5242c5fbabd149c8812c8152ab2e64929c4778e8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 18:28:28 -0700 Subject: [PATCH 26/28] test - auth exception handler --- .../proxy/auth/test_auth_exception_handler.py | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/tests/litellm/proxy/auth/test_auth_exception_handler.py b/tests/litellm/proxy/auth/test_auth_exception_handler.py index ca0e8c88c4..b44de86e05 100644 --- a/tests/litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/litellm/proxy/auth/test_auth_exception_handler.py @@ -101,9 +101,33 @@ def test_should_allow_request_on_db_unavailable_false(): assert handler.should_allow_request_on_db_unavailable() == False -# Test _handle_authentication_error method @pytest.mark.asyncio -async def test_handle_authentication_error_db_unavailable(): +@pytest.mark.parametrize( + "prisma_error", + [ + PrismaError(), + DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + ForeignKeyViolationError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + MissingRequiredValueError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + TableNotFoundError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + RecordNotFoundError( + data={"user_facing_error": {"meta": {"table": "test_table"}}} + ), + HTTPClientClosedError(), + ClientNotConnectedError(), + ], +) +async def test_handle_authentication_error_db_unavailable(prisma_error): handler = UserAPIKeyAuthExceptionHandler() # Mock request and other dependencies @@ -118,9 +142,8 @@ async def test_handle_authentication_error_db_unavailable(): "litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": True}, ): - db_error = prisma_errors.PrismaError() result = await handler._handle_authentication_error( - db_error, + prisma_error, mock_request, mock_request_data, mock_route, From 763f853a9f9f6e678d966e08df91ccec3a76009d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 18:32:58 -0700 Subject: [PATCH 27/28] docs fix --- docs/my-website/docs/providers/bedrock.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 4f88fdb39b..ed65e14b8b 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -1776,6 +1776,7 @@ response = completion( ) ``` + 1. Setup config.yaml @@ -1820,11 +1821,13 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \ ``` + ### SSO Login (AWS Profile) - Set `AWS_PROFILE` environment variable - Make bedrock completion call + ```python import os from litellm import completion @@ -1940,9 +1943,6 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/images/generations' \ "colorGuidedGenerationParams":{"colors":["#FFFFFF"]} }' ``` - - - | Model Name | Function Call | |-------------------------|---------------------------------------------| From 485aa87e6505e3749d0a15a0c335233af5c02136 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 26 Mar 2025 18:48:18 -0700 Subject: [PATCH 28/28] allow_requests_on_db_unavailable --- docs/my-website/docs/proxy/config_settings.md | 2 +- docs/my-website/docs/proxy/prod.md | 110 +++--------------- 2 files changed, 18 insertions(+), 94 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 0093464d93..4a62184df7 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -160,7 +160,7 @@ general_settings: | database_url | string | The URL for the database connection [Set up Virtual Keys](virtual_keys) | | database_connection_pool_limit | integer | The limit for database connection pool [Setting DB Connection Pool limit](#configure-db-pool-limits--connection-timeouts) | | database_connection_timeout | integer | The timeout for database connections in seconds [Setting DB Connection Pool limit, timeout](#configure-db-pool-limits--connection-timeouts) | -| allow_requests_on_db_unavailable | boolean | If true, allows requests to succeed even if DB is unreachable. **Only use this if running LiteLLM in your VPC** This will allow requests to work even when LiteLLM cannot connect to the DB to verify a Virtual Key | +| allow_requests_on_db_unavailable | boolean | If true, allows requests to succeed even if DB is unreachable. **Only use this if running LiteLLM in your VPC** This will allow requests to work even when LiteLLM cannot connect to the DB to verify a Virtual Key [Doc on graceful db unavailability](prod#5-if-running-litellm-on-vpc-gracefully-handle-db-unavailability) | | custom_auth | string | Write your own custom authentication logic [Doc Custom Auth](virtual_keys#custom-auth) | | max_parallel_requests | integer | The max parallel requests allowed per deployment | | global_max_parallel_requests | integer | The max parallel requests allowed on the proxy overall | diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index d3ba2d6224..314300f2a0 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -94,15 +94,29 @@ This disables the load_dotenv() functionality, which will automatically load you ## 5. If running LiteLLM on VPC, gracefully handle DB unavailability -This will allow LiteLLM to continue to process requests even if the DB is unavailable. This is better handling for DB unavailability. +When running LiteLLM on a VPC (and inaccessible from the public internet), you can enable graceful degradation so that request processing continues even if the database is temporarily unavailable. + **WARNING: Only do this if you're running LiteLLM on VPC, that cannot be accessed from the public internet.** -```yaml +#### Configuration + +```yaml showLineNumbers title="litellm config.yaml" general_settings: allow_requests_on_db_unavailable: True ``` +#### Expected Behavior + +When `allow_requests_on_db_unavailable` is set to `true`, LiteLLM will handle errors as follows: + +| Type of Error | Expected Behavior | Details | +|---------------|-------------------|----------------| +| Prisma Errors | ✅ Request will be allowed | Covers issues like DB connection resets or rejections from the DB via Prisma, the ORM used by LiteLLM. | +| Httpx Errors | ✅ Request will be allowed | Occurs when the database is unreachable, allowing the request to proceed despite the DB outage. | +| LiteLLM Budget Errors or Model Errors | ❌ Request will be blocked | Triggered when the DB is reachable but the authentication token is invalid, lacks access, or exceeds budget limits. | + + ## 6. Disable spend_logs & error_logs if not using the LiteLLM UI By default, LiteLLM writes several types of logs to the database: @@ -182,94 +196,4 @@ You should only see the following level of details in logs on the proxy server # INFO: 192.168.2.205:11774 - "POST /chat/completions HTTP/1.1" 200 OK # INFO: 192.168.2.205:34717 - "POST /chat/completions HTTP/1.1" 200 OK # INFO: 192.168.2.205:29734 - "POST /chat/completions HTTP/1.1" 200 OK -``` - - -### Machine Specifications to Deploy LiteLLM - -| Service | Spec | CPUs | Memory | Architecture | Version| -| --- | --- | --- | --- | --- | --- | -| Server | `t2.small`. | `1vCPUs` | `8GB` | `x86` | -| Redis Cache | - | - | - | - | 7.0+ Redis Engine| - - -### Reference Kubernetes Deployment YAML - -Reference Kubernetes `deployment.yaml` that was load tested by us - -```yaml -apiVersion: apps/v1 -kind: Deployment -metadata: - name: litellm-deployment -spec: - replicas: 3 - selector: - matchLabels: - app: litellm - template: - metadata: - labels: - app: litellm - spec: - containers: - - name: litellm-container - image: ghcr.io/berriai/litellm:main-latest - imagePullPolicy: Always - env: - - name: AZURE_API_KEY - value: "d6******" - - name: AZURE_API_BASE - value: "https://ope******" - - name: LITELLM_MASTER_KEY - value: "sk-1234" - - name: DATABASE_URL - value: "po**********" - args: - - "--config" - - "/app/proxy_config.yaml" # Update the path to mount the config file - volumeMounts: # Define volume mount for proxy_config.yaml - - name: config-volume - mountPath: /app - readOnly: true - livenessProbe: - httpGet: - path: /health/liveliness - port: 4000 - initialDelaySeconds: 120 - periodSeconds: 15 - successThreshold: 1 - failureThreshold: 3 - timeoutSeconds: 10 - readinessProbe: - httpGet: - path: /health/readiness - port: 4000 - initialDelaySeconds: 120 - periodSeconds: 15 - successThreshold: 1 - failureThreshold: 3 - timeoutSeconds: 10 - volumes: # Define volume to mount proxy_config.yaml - - name: config-volume - configMap: - name: litellm-config - -``` - - -Reference Kubernetes `service.yaml` that was load tested by us -```yaml -apiVersion: v1 -kind: Service -metadata: - name: litellm-service -spec: - selector: - app: litellm - ports: - - protocol: TCP - port: 4000 - targetPort: 4000 - type: LoadBalancer -``` +``` \ No newline at end of file