github/OpenHands
1
0
Fork 0
mirror of https://github.com/OpenHands/OpenHands.git synced 2026-03-22 05:37:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

PLTF-309: disable budget enforcement when ENABLE_BILLING=false (#13440)

Browse Source 
Co-authored-by: openhands <openhands@all-hands.dev>
This commit is contained in:
aivong-openhands
2026-03-17 14:26:13 -05:00
committed by GitHub
parent 09ca1b882f
commit 855ef7ba5f
2 changed files with 101 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
enterprise/storage/lite_llm_manager.py
View File

@@ -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...

39
enterprise/tests/unit/test_lite_llm_manager.py
View File

@@ -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')