Skip to content

Commit

Permalink
(fix) proxy - fix when STORE_MODEL_IN_DB should be set (BerriAI#6492)
Browse files Browse the repository at this point in the history
* set store_model_in_db at the top

* correctly use store_model_in_db global
  • Loading branch information
ishaan-jaff authored Oct 29, 2024
1 parent 441adad commit 8e19a31
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 5 deletions.
15 changes: 10 additions & 5 deletions litellm/proxy/proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,12 @@ def generate_feedback_box():
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,
Expand Down Expand Up @@ -2894,9 +2899,9 @@ async def initialize_scheduled_background_jobs(
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
Expand All @@ -2921,8 +2926,9 @@ async def initialize_scheduled_background_jobs(

### 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,
Expand Down Expand Up @@ -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,
)


Expand Down
23 changes: 23 additions & 0 deletions litellm/secret_managers/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 8e19a31

Please sign in to comment.