From b0db75df1fb9f5871cdea662035ee719ef0a9149 Mon Sep 17 00:00:00 2001 From: rstar327 Date: Tue, 17 Mar 2026 13:35:07 -0400 Subject: [PATCH] fix(proxy): convert max_budget to float when set from environment variable (#23855) Fixes #23843 --- litellm/proxy/proxy_server.py | 68 +++++++++---------- .../proxy/test_max_budget_env_var.py | 38 +++++++++++ 2 files changed, 71 insertions(+), 35 deletions(-) create mode 100644 tests/test_litellm/proxy/test_max_budget_env_var.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9c29927c5cb..a775b12b206 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -638,9 +638,9 @@ except ImportError: server_root_path = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1523,9 +1523,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1533,13 +1533,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[ - str -] = [] # Models that should use native provider background mode instead of polling +native_background_mode: List[str] = ( + [] +) # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -1898,9 +1898,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -2021,11 +2021,9 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug( - f""" + verbose_proxy_logger.debug(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """ - ) + """) def _get_process_rss_mb() -> Optional[float]: @@ -3303,7 +3301,7 @@ class ProxyConfig: async_only_mode=True # only init async clients ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid - ) # type:ignore + ) # type: ignore if redis_usage_cache is not None and router.cache.redis_cache is None: router._update_redis_cache(cache=redis_usage_cache) @@ -4952,10 +4950,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -5337,9 +5335,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -5357,8 +5355,8 @@ async def initialize( # noqa: PLR0915 litellm.add_function_to_prompt = True dynamic_config["general"]["add_function_to_prompt"] = True if max_budget: # litellm-specific param - litellm.max_budget = max_budget - dynamic_config["general"]["max_budget"] = max_budget + litellm.max_budget = float(max_budget) + dynamic_config["general"]["max_budget"] = litellm.max_budget if experimental: pass user_telemetry = telemetry @@ -5676,9 +5674,9 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[ - Union[bool, str] - ] = general_settings.get("use_redis_transaction_buffer", False) + _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( + general_settings.get("use_redis_transaction_buffer", False) + ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -12114,9 +12112,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/tests/test_litellm/proxy/test_max_budget_env_var.py b/tests/test_litellm/proxy/test_max_budget_env_var.py new file mode 100644 index 00000000000..4b965392823 --- /dev/null +++ b/tests/test_litellm/proxy/test_max_budget_env_var.py @@ -0,0 +1,38 @@ +""" +Test that max_budget from environment variable (string) is correctly +converted to float. +GitHub Issue: #23843 +""" + +import pytest + +import litellm +from litellm.proxy.proxy_server import initialize + + +@pytest.mark.asyncio +async def test_max_budget_string_converted_to_float(): + """ + When max_budget is set via os.environ/MAX_BUDGET, it arrives as a + string. initialize() should convert it to float so the comparison + `litellm.max_budget > 0` doesn't raise TypeError. + """ + original = litellm.max_budget + try: + await initialize(max_budget="100.5") + assert isinstance(litellm.max_budget, float) + assert litellm.max_budget == 100.5 + finally: + litellm.max_budget = original + + +@pytest.mark.asyncio +async def test_max_budget_float_stays_float(): + """max_budget as float should still work.""" + original = litellm.max_budget + try: + await initialize(max_budget=200.0) + assert isinstance(litellm.max_budget, float) + assert litellm.max_budget == 200.0 + finally: + litellm.max_budget = original