diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cadb6063c..1f0271f34 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -250,7 +250,12 @@ from litellm.secret_managers.aws_secret_manager import ( load_aws_secret_manager, ) from litellm.secret_managers.google_kms import load_google_kms -from litellm.secret_managers.main import get_secret, get_secret_str, str_to_bool +from litellm.secret_managers.main import ( + get_secret, + get_secret_bool, + get_secret_str, + str_to_bool, +) from litellm.types.integrations.slack_alerting import SlackAlertingArgs from litellm.types.llms.anthropic import ( AnthropicMessagesRequest, @@ -2894,9 +2899,9 @@ class ProxyStartupEvent: proxy_budget_rescheduler_max_time: int, proxy_batch_write_at: int, proxy_logging_obj: ProxyLogging, - store_model_in_db: bool, ): """Initializes scheduled background jobs""" + global store_model_in_db scheduler = AsyncIOScheduler() interval = random.randint( proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time @@ -2921,8 +2926,9 @@ class ProxyStartupEvent: ### ADD NEW MODELS ### store_model_in_db = ( - get_secret("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db - ) # type: ignore + get_secret_bool("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db + ) + if store_model_in_db is True: scheduler.add_job( proxy_config.add_deployment, @@ -3141,7 +3147,6 @@ async def startup_event(): proxy_budget_rescheduler_max_time=proxy_budget_rescheduler_max_time, proxy_batch_write_at=proxy_batch_write_at, proxy_logging_obj=proxy_logging_obj, - store_model_in_db=store_model_in_db, ) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 522f2bc39..f3d6d420a 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -67,6 +67,29 @@ def get_secret_str( return value +def get_secret_bool( + secret_name: str, + default_value: Optional[bool] = None, +) -> Optional[bool]: + """ + Guarantees response from 'get_secret' is either boolean or none. Used for fixing linting errors. + + Args: + secret_name: The name of the secret to get. + default_value: The default value to return if the secret is not found. + + Returns: + The secret value as a boolean or None if the secret is not found. + """ + _secret_value = get_secret(secret_name, default_value) + if _secret_value is None: + return None + elif isinstance(_secret_value, bool): + return _secret_value + else: + return str_to_bool(_secret_value) + + def get_secret( # noqa: PLR0915 secret_name: str, default_value: Optional[Union[str, bool]] = None,