diff --git a/litellm/__init__.py b/litellm/__init__.py index 3f2a1e4b4d..26b761c64a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -146,6 +146,7 @@ suppress_debug_info = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None +upperbound_key_generate_params: Optional[Dict] = None default_team_settings: Optional[List] = None #### RELIABILITY #### request_timeout: Optional[float] = 6000 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 289a36cb2b..494c874147 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1391,6 +1391,26 @@ class ProxyConfig: proxy_config = ProxyConfig() +def _duration_in_seconds(duration: str): + match = re.match(r"(\d+)([smhd]?)", duration) + if not match: + raise ValueError("Invalid duration format") + + value, unit = match.groups() + value = int(value) + + if unit == "s": + return value + elif unit == "m": + return value * 60 + elif unit == "h": + return value * 3600 + elif unit == "d": + return value * 86400 + else: + raise ValueError("Unsupported duration unit") + + async def generate_key_helper_fn( duration: Optional[str], models: list, @@ -1425,25 +1445,6 @@ async def generate_key_helper_fn( if token is None: token = f"sk-{secrets.token_urlsafe(16)}" - def _duration_in_seconds(duration: str): - match = re.match(r"(\d+)([smhd]?)", duration) - if not match: - raise ValueError("Invalid duration format") - - value, unit = match.groups() - value = int(value) - - if unit == "s": - return value - elif unit == "m": - return value * 60 - elif unit == "h": - return value * 3600 - elif unit == "d": - return value * 86400 - else: - raise ValueError("Unsupported duration unit") - if duration is None: # allow tokens that never expire expires = None else: @@ -2660,6 +2661,36 @@ async def generate_key_fn( elif key == "metadata" and value == {}: setattr(data, key, litellm.default_key_generate_params.get(key, {})) + # check if user set default key/generate params on config.yaml + if litellm.upperbound_key_generate_params is not None: + for elem in data: + # if key in litellm.upperbound_key_generate_params, use the min of value and litellm.upperbound_key_generate_params[key] + key, value = elem + if value is not None and key in litellm.upperbound_key_generate_params: + # if value is float/int + if key in [ + "max_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + ]: + if value > litellm.upperbound_key_generate_params[key]: + # directly compare floats/ints + setattr( + data, key, litellm.upperbound_key_generate_params[key] + ) + elif key == "budget_duration": + # budgets are in 1s, 1m, 1h, 1d, 1m (30s, 30m, 30h, 30d, 30m) + # compare the duration in seconds and max duration in seconds + upperbound_budget_duration = _duration_in_seconds( + duration=litellm.upperbound_key_generate_params[key] + ) + user_set_budget_duration = _duration_in_seconds(duration=value) + if user_set_budget_duration > upperbound_budget_duration: + setattr( + data, key, litellm.upperbound_key_generate_params[key] + ) + data_json = data.json() # type: ignore # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users