diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py index 836ebe8278..725b8147a3 100644 --- a/enterprise/storage/lite_llm_manager.py +++ b/enterprise/storage/lite_llm_manager.py @@ -29,14 +29,37 @@ KEY_VERIFICATION_TIMEOUT = 5.0 # A very large number to represent "unlimited" until LiteLLM fixes their unlimited update bug. UNLIMITED_BUDGET_SETTING = 1000000000.0 -try: - DEFAULT_INITIAL_BUDGET = float(os.environ.get('DEFAULT_INITIAL_BUDGET', 0.0)) - if DEFAULT_INITIAL_BUDGET < 0: +# Check if billing is enabled (defaults to false for enterprise deployments) +ENABLE_BILLING = os.environ.get('ENABLE_BILLING', 'false').lower() == 'true' + + +def _get_default_initial_budget() -> float | None: + """Get the default initial budget for new teams. + + When billing is disabled (ENABLE_BILLING=false), returns None to disable + budget enforcement in LiteLLM. When billing is enabled, returns the + DEFAULT_INITIAL_BUDGET environment variable value (default 0.0). + + Returns: + float | None: The default budget, or None to disable budget enforcement. + """ + if not ENABLE_BILLING: + return None + + try: + budget = float(os.environ.get('DEFAULT_INITIAL_BUDGET', 0.0)) + if budget < 0: + raise ValueError( + f'DEFAULT_INITIAL_BUDGET must be non-negative, got {budget}' + ) + return budget + except ValueError as e: raise ValueError( - f'DEFAULT_INITIAL_BUDGET must be non-negative, got {DEFAULT_INITIAL_BUDGET}' - ) -except ValueError as e: - raise ValueError(f'Invalid DEFAULT_INITIAL_BUDGET environment variable: {e}') from e + f'Invalid DEFAULT_INITIAL_BUDGET environment variable: {e}' + ) from e + + +DEFAULT_INITIAL_BUDGET: float | None = _get_default_initial_budget() def get_openhands_cloud_key_alias(keycloak_user_id: str, org_id: str) -> str: @@ -110,12 +133,15 @@ class LiteLlmManager: ) as client: # Check if team already exists and get its budget # New users joining existing orgs should inherit the team's budget - team_budget: float = DEFAULT_INITIAL_BUDGET + # When billing is disabled, DEFAULT_INITIAL_BUDGET is None + team_budget: float | None = DEFAULT_INITIAL_BUDGET try: existing_team = await LiteLlmManager._get_team(client, org_id) if existing_team: team_info = existing_team.get('team_info', {}) - team_budget = team_info.get('max_budget', 0.0) or 0.0 + # Preserve None from existing team (no budget enforcement) + existing_budget = team_info.get('max_budget') + team_budget = existing_budget logger.info( 'LiteLlmManager:create_entries:existing_team_budget', extra={ @@ -525,8 +551,17 @@ class LiteLlmManager: client: httpx.AsyncClient, team_alias: str, team_id: str, - max_budget: float, + max_budget: float | None, ): + """Create a new team in LiteLLM. + + Args: + client: The HTTP client to use. + team_alias: The alias for the team. + team_id: The ID for the team. + max_budget: The maximum budget for the team. When None, budget + enforcement is disabled (unlimited usage). + """ if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: logger.warning('LiteLLM API configuration not found') return @@ -536,7 +571,7 @@ class LiteLlmManager: 'team_id': team_id, 'team_alias': team_alias, 'models': [], - 'max_budget': max_budget, + 'max_budget': max_budget, # None disables budget enforcement 'spend': 0, 'metadata': { 'version': ORG_SETTINGS_VERSION, @@ -918,8 +953,17 @@ class LiteLlmManager: client: httpx.AsyncClient, keycloak_user_id: str, team_id: str, - max_budget: float, + max_budget: float | None, ): + """Add a user to a team in LiteLLM. + + Args: + client: The HTTP client to use. + keycloak_user_id: The user's Keycloak ID. + team_id: The team ID. + max_budget: The maximum budget for the user in the team. When None, + budget enforcement is disabled (unlimited usage). + """ if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: logger.warning('LiteLLM API configuration not found') return @@ -928,7 +972,7 @@ class LiteLlmManager: json={ 'team_id': team_id, 'member': {'user_id': keycloak_user_id, 'role': 'user'}, - 'max_budget_in_team': max_budget, + 'max_budget_in_team': max_budget, # None disables budget enforcement }, ) # Failed to add user to team - this is an unforseen error state... @@ -998,8 +1042,17 @@ class LiteLlmManager: client: httpx.AsyncClient, keycloak_user_id: str, team_id: str, - max_budget: float, + max_budget: float | None, ): + """Update a user's budget in a team. + + Args: + client: The HTTP client to use. + keycloak_user_id: The user's Keycloak ID. + team_id: The team ID. + max_budget: The maximum budget for the user in the team. When None, + budget enforcement is disabled (unlimited usage). + """ if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: logger.warning('LiteLLM API configuration not found') return @@ -1008,7 +1061,7 @@ class LiteLlmManager: json={ 'team_id': team_id, 'user_id': keycloak_user_id, - 'max_budget_in_team': max_budget, + 'max_budget_in_team': max_budget, # None disables budget enforcement }, ) # Failed to update user in team - this is an unforseen error state... diff --git a/enterprise/tests/unit/test_lite_llm_manager.py b/enterprise/tests/unit/test_lite_llm_manager.py index 1f7623d79c..0cfc9fe58b 100644 --- a/enterprise/tests/unit/test_lite_llm_manager.py +++ b/enterprise/tests/unit/test_lite_llm_manager.py @@ -38,8 +38,9 @@ class TestDefaultInitialBudget: if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] - # Clear the env var + # Clear the env vars os.environ.pop('DEFAULT_INITIAL_BUDGET', None) + os.environ.pop('ENABLE_BILLING', None) # Restore original module or reimport fresh if original_module is not None: @@ -47,31 +48,56 @@ class TestDefaultInitialBudget: else: importlib.import_module('storage.lite_llm_manager') - def test_default_initial_budget_defaults_to_zero(self): - """Test that DEFAULT_INITIAL_BUDGET defaults to 0.0 when env var not set.""" + def test_default_initial_budget_none_when_billing_disabled(self): + """Test that DEFAULT_INITIAL_BUDGET is None when billing is disabled.""" # Temporarily remove the module so we can reimport with different env vars if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] - # Clear the env var and reimport + # Ensure billing is disabled (default) and reimport + os.environ.pop('ENABLE_BILLING', None) + os.environ.pop('DEFAULT_INITIAL_BUDGET', None) + module = importlib.import_module('storage.lite_llm_manager') + assert module.DEFAULT_INITIAL_BUDGET is None + + def test_default_initial_budget_defaults_to_zero_when_billing_enabled(self): + """Test that DEFAULT_INITIAL_BUDGET defaults to 0.0 when billing is enabled.""" + # Temporarily remove the module so we can reimport with different env vars + if 'storage.lite_llm_manager' in sys.modules: + del sys.modules['storage.lite_llm_manager'] + + # Enable billing and reimport + os.environ['ENABLE_BILLING'] = 'true' os.environ.pop('DEFAULT_INITIAL_BUDGET', None) module = importlib.import_module('storage.lite_llm_manager') assert module.DEFAULT_INITIAL_BUDGET == 0.0 - def test_default_initial_budget_uses_env_var(self): - """Test that DEFAULT_INITIAL_BUDGET uses value from environment variable.""" + def test_default_initial_budget_uses_env_var_when_billing_enabled(self): + """Test that DEFAULT_INITIAL_BUDGET uses value from environment variable when billing enabled.""" if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] + os.environ['ENABLE_BILLING'] = 'true' os.environ['DEFAULT_INITIAL_BUDGET'] = '100.0' module = importlib.import_module('storage.lite_llm_manager') assert module.DEFAULT_INITIAL_BUDGET == 100.0 + def test_default_initial_budget_ignores_env_var_when_billing_disabled(self): + """Test that DEFAULT_INITIAL_BUDGET returns None when billing disabled, ignoring env var.""" + if 'storage.lite_llm_manager' in sys.modules: + del sys.modules['storage.lite_llm_manager'] + + os.environ.pop('ENABLE_BILLING', None) # billing disabled by default + os.environ['DEFAULT_INITIAL_BUDGET'] = '100.0' + module = importlib.import_module('storage.lite_llm_manager') + assert module.DEFAULT_INITIAL_BUDGET is None + def test_default_initial_budget_rejects_invalid_value(self): """Test that DEFAULT_INITIAL_BUDGET raises ValueError for invalid values.""" if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] + os.environ['ENABLE_BILLING'] = 'true' os.environ['DEFAULT_INITIAL_BUDGET'] = 'abc' with pytest.raises(ValueError) as exc_info: importlib.import_module('storage.lite_llm_manager') @@ -82,6 +108,7 @@ class TestDefaultInitialBudget: if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] + os.environ['ENABLE_BILLING'] = 'true' os.environ['DEFAULT_INITIAL_BUDGET'] = '-10.0' with pytest.raises(ValueError) as exc_info: importlib.import_module('storage.lite_llm_manager')