From e50e3b5ce7cfabbf78418992b791b6ca7fd7e5f6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Nov 2024 12:58:03 -0800 Subject: [PATCH 1/6] fix use helper for _handle_failed_db_connection_for_get_key_object --- litellm/proxy/auth/auth_checks.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e00d494d94b..e8b27966fbb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,6 +13,7 @@ import traceback from datetime import datetime from typing import TYPE_CHECKING, Any, List, Literal, Optional +import httpx from pydantic import BaseModel import litellm @@ -717,12 +718,36 @@ async def get_key_object( ) return _response - except Exception: + except httpx.ConnectError as e: + return _handle_failed_db_connection_for_get_key_object(e=e) + except Exception as e: raise Exception( f"Key doesn't exist in db. key={hashed_token}. Create key via `/key/generate` call." ) +def _handle_failed_db_connection_for_get_key_object(e: Exception): + """ + 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_failed_db_requests is True + + Raises: + - Orignal Exception in all other cases + """ + from litellm.proxy.proxy_server import general_settings + + if general_settings.get("allow_failed_db_requests", True): + return UserAPIKeyAuth( + key_name="failed-to-connect-to-db", token="failed-to-connect-to-db" + ) + else: + raise e + + @log_to_opentelemetry async def get_org_object( org_id: str, From 8d7934c4ba77449a9bb6b38bb0d6e17c79b60030 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Nov 2024 13:12:14 -0800 Subject: [PATCH 2/6] track ALLOW_FAILED_DB_REQUESTS on prometheus --- litellm/proxy/auth/auth_checks.py | 25 ++++++++++++++++++++++--- litellm/proxy/proxy_config.yaml | 6 +++--- litellm/types/services.py | 1 + 3 files changed, 26 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e8b27966fbb..9e535a3eb2a 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -719,14 +719,16 @@ async def get_key_object( return _response except httpx.ConnectError as e: - return _handle_failed_db_connection_for_get_key_object(e=e) + return await _handle_failed_db_connection_for_get_key_object(e=e) except Exception as e: raise Exception( f"Key doesn't exist in db. key={hashed_token}. Create key via `/key/generate` call." ) -def _handle_failed_db_connection_for_get_key_object(e: Exception): +async def _handle_failed_db_connection_for_get_key_object( + e: Exception, +) -> UserAPIKeyAuth: """ Handles httpx.ConnectError when reading a Virtual Key from LiteLLM DB @@ -738,9 +740,26 @@ def _handle_failed_db_connection_for_get_key_object(e: Exception): Raises: - Orignal Exception in all other cases """ - from litellm.proxy.proxy_server import general_settings + from litellm.proxy.proxy_server import general_settings, proxy_logging_obj + # If this flag is on, requests failing to connect to the DB will be allowed if general_settings.get("allow_failed_db_requests", True): + # log to prometheus + await proxy_logging_obj.service_logging_obj.async_service_failure_hook( + error=e, + service=ServiceTypes.ALLOW_FAILED_DB_REQUESTS, + call_type="get_key_object", + parent_otel_span=None, + duration=0.0, + start_time=None, + end_time=None, + event_metadata={ + "function_name": "get_key_object", + "function_kwargs": {}, + "function_args": [], + }, + ) + return UserAPIKeyAuth( key_name="failed-to-connect-to-db", token="failed-to-connect-to-db" ) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 23834f75970..d79bcdb0341 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,6 +6,6 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ -general_settings: - alerting: ["slack"] - alerting_threshold: 0.001 +litellm_settings: + callbacks: ["prometheus"] + service_callback: ["prometheus_system"] diff --git a/litellm/types/services.py b/litellm/types/services.py index cfa427ebc3d..161832404c1 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -12,6 +12,7 @@ class ServiceTypes(str, enum.Enum): REDIS = "redis" DB = "postgres" + ALLOW_FAILED_DB_REQUESTS = "allow_failed_db_requests" BATCH_WRITE_TO_DB = "batch_write_to_db" LITELLM = "self" ROUTER = "router" From 9b2861f11f0fe5d2af3224666db40091c821e0d3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Nov 2024 13:15:37 -0800 Subject: [PATCH 3/6] fix allow_failed_db_requests check --- 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 9e535a3eb2a..50e22654a42 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -743,7 +743,7 @@ async def _handle_failed_db_connection_for_get_key_object( from litellm.proxy.proxy_server import general_settings, proxy_logging_obj # If this flag is on, requests failing to connect to the DB will be allowed - if general_settings.get("allow_failed_db_requests", True): + if general_settings.get("allow_failed_db_requests", False) is True: # log to prometheus await proxy_logging_obj.service_logging_obj.async_service_failure_hook( error=e, From 7582266267ebcd62e240cbb9423426e2159e65b3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Nov 2024 13:19:06 -0800 Subject: [PATCH 4/6] fix allow_requests_on_db_unavailable --- litellm/proxy/auth/auth_checks.py | 14 ++++---------- litellm/types/services.py | 2 +- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 50e22654a42..d2b04122b43 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -735,7 +735,7 @@ async def _handle_failed_db_connection_for_get_key_object( Use this if you don't want failed DB queries to block LLM API reqiests Returns: - - UserAPIKeyAuth: If general_settings.allow_failed_db_requests is True + - UserAPIKeyAuth: If general_settings.allow_requests_on_db_unavailable is True Raises: - Orignal Exception in all other cases @@ -743,21 +743,15 @@ async def _handle_failed_db_connection_for_get_key_object( from litellm.proxy.proxy_server import general_settings, proxy_logging_obj # If this flag is on, requests failing to connect to the DB will be allowed - if general_settings.get("allow_failed_db_requests", False) is True: + if general_settings.get("allow_requests_on_db_unavailable", False) is True: # log to prometheus - await proxy_logging_obj.service_logging_obj.async_service_failure_hook( - error=e, - service=ServiceTypes.ALLOW_FAILED_DB_REQUESTS, + proxy_logging_obj.service_logging_obj.service_success_hook( + service=ServiceTypes.ALLOW_REQUESTS_ON_DB_UNAVAILABLE, call_type="get_key_object", parent_otel_span=None, duration=0.0, start_time=None, end_time=None, - event_metadata={ - "function_name": "get_key_object", - "function_kwargs": {}, - "function_args": [], - }, ) return UserAPIKeyAuth( diff --git a/litellm/types/services.py b/litellm/types/services.py index 161832404c1..6ffdfc8916f 100644 --- a/litellm/types/services.py +++ b/litellm/types/services.py @@ -12,7 +12,7 @@ class ServiceTypes(str, enum.Enum): REDIS = "redis" DB = "postgres" - ALLOW_FAILED_DB_REQUESTS = "allow_failed_db_requests" + ALLOW_REQUESTS_ON_DB_UNAVAILABLE = "allow_requests_on_db_unavailable" BATCH_WRITE_TO_DB = "batch_write_to_db" LITELLM = "self" ROUTER = "router" From 761d56e808ff674fe37c76925498db966cc94420 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Nov 2024 13:25:46 -0800 Subject: [PATCH 5/6] fix allow_requests_on_db_unavailable --- docs/my-website/docs/proxy/configs.md | 5 +++++ litellm/proxy/proxy_config.yaml | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index 28b0b67e336..265670327b6 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -692,9 +692,13 @@ general_settings: allowed_routes: ["route1", "route2"] # list of allowed proxy API routes - a user can access. (currently JWT-Auth only) key_management_system: google_kms # either google_kms or azure_kms master_key: string + + # Database Settings database_url: string database_connection_pool_limit: 0 # default 100 database_connection_timeout: 0 # default 60s + allow_requests_on_db_unavailable: boolean # if true, will allow requests that can not connect to the DB to verify Virtual Key to still work + custom_auth: string max_parallel_requests: 0 # the max parallel requests allowed per deployment global_max_parallel_requests: 0 # the max parallel requests allowed on the proxy all up @@ -766,6 +770,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** | | 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/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index d79bcdb0341..6635fb2ff43 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -9,3 +9,7 @@ model_list: litellm_settings: callbacks: ["prometheus"] service_callback: ["prometheus_system"] + + +general_settings: + allow_requests_on_db_unavailable: true From d0521d693ef060b1efe217df63965208d19aee91 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Nov 2024 13:27:24 -0800 Subject: [PATCH 6/6] docs allow_requests_on_db_unavailable --- docs/my-website/docs/proxy/configs.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index 265670327b6..1adc4943d33 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -770,7 +770,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** | +| 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 | | 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 |