phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(fix) proxy - fix when STORE_MODEL_IN_DB should be set (#6492)

Browse source 
* set store_model_in_db at the top

* correctly use store_model_in_db global
This commit is contained in:
Ishaan Jaff 2024-10-29 21:28:14 +05:30 • committed by GitHub
parent 441adad3ae
commit 8e19a31d36
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 33 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

15
litellm/proxy/proxy_server.py
View file

@ -250,7 +250,12 @@ from litellm.secret_managers.aws_secret_manager import (
load_aws_secret_manager, load_aws_secret_manager,
) )
from litellm.secret_managers.google_kms import load_google_kms 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.integrations.slack_alerting import SlackAlertingArgs
from litellm.types.llms.anthropic import ( from litellm.types.llms.anthropic import (
AnthropicMessagesRequest, AnthropicMessagesRequest,
@ -2894,9 +2899,9 @@ class ProxyStartupEvent:
proxy_budget_rescheduler_max_time: int, proxy_budget_rescheduler_max_time: int,
proxy_batch_write_at: int, proxy_batch_write_at: int,
proxy_logging_obj: ProxyLogging, proxy_logging_obj: ProxyLogging,
store_model_in_db: bool,
): ):
"""Initializes scheduled background jobs""" """Initializes scheduled background jobs"""
global store_model_in_db
scheduler = AsyncIOScheduler() scheduler = AsyncIOScheduler()
interval = random.randint( interval = random.randint(
proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time
@ -2921,8 +2926,9 @@ class ProxyStartupEvent:
### ADD NEW MODELS ### ### ADD NEW MODELS ###
store_model_in_db = ( store_model_in_db = (
get_secret("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db get_secret_bool("STORE_MODEL_IN_DB", store_model_in_db) or store_model_in_db
) # type: ignore )
if store_model_in_db is True: if store_model_in_db is True:
scheduler.add_job( scheduler.add_job(
proxy_config.add_deployment, proxy_config.add_deployment,
@ -3141,7 +3147,6 @@ async def startup_event():
proxy_budget_rescheduler_max_time=proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_max_time=proxy_budget_rescheduler_max_time,
proxy_batch_write_at=proxy_batch_write_at, proxy_batch_write_at=proxy_batch_write_at,
proxy_logging_obj=proxy_logging_obj, proxy_logging_obj=proxy_logging_obj,
store_model_in_db=store_model_in_db,
) )

23
litellm/secret_managers/main.py
View file

@ -67,6 +67,29 @@ def get_secret_str(
return value 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 def get_secret( # noqa: PLR0915
secret_name: str, secret_name: str,
default_value: Optional[Union[str, bool]] = None, default_value: Optional[Union[str, bool]] = None,