Fix when budgets are None

This commit is contained in:
Chuck Butkus
2026-03-18 21:50:47 -04:00
parent 3c5023ac02
commit 417f56a5b6

View File

@@ -530,20 +530,26 @@ class LiteLlmManager:
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
logger.warning('LiteLLM API configuration not found') logger.warning('LiteLLM API configuration not found')
return return
json_data: dict[str, Any] = {
'team_id': team_id,
'team_alias': team_alias,
'models': [],
'spend': 0,
'metadata': {
'version': ORG_SETTINGS_VERSION,
'model': get_default_litellm_model(),
},
}
if max_budget is not None:
json_data['max_budget'] = max_budget
response = await client.post( response = await client.post(
f'{LITE_LLM_API_URL}/team/new', f'{LITE_LLM_API_URL}/team/new',
json={ json=json_data,
'team_id': team_id,
'team_alias': team_alias,
'models': [],
'max_budget': max_budget,
'spend': 0,
'metadata': {
'version': ORG_SETTINGS_VERSION,
'model': get_default_litellm_model(),
},
},
) )
# Team failed to create in litellm - this is an unforseen error state... # Team failed to create in litellm - this is an unforseen error state...
if not response.is_success: if not response.is_success:
if ( if (
@@ -923,14 +929,20 @@ class LiteLlmManager:
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
logger.warning('LiteLLM API configuration not found') logger.warning('LiteLLM API configuration not found')
return return
json_data: dict[str, Any] = {
'team_id': team_id,
'member': {'user_id': keycloak_user_id, 'role': 'user'},
}
if max_budget is not None:
json_data['max_budget_in_team'] = max_budget
response = await client.post( response = await client.post(
f'{LITE_LLM_API_URL}/team/member_add', f'{LITE_LLM_API_URL}/team/member_add',
json={ json=json_data,
'team_id': team_id,
'member': {'user_id': keycloak_user_id, 'role': 'user'},
'max_budget_in_team': max_budget,
},
) )
# Failed to add user to team - this is an unforseen error state... # Failed to add user to team - this is an unforseen error state...
if not response.is_success: if not response.is_success:
if ( if (
@@ -1003,14 +1015,20 @@ class LiteLlmManager:
if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None:
logger.warning('LiteLLM API configuration not found') logger.warning('LiteLLM API configuration not found')
return return
json_data: dict[str, Any] = {
'team_id': team_id,
'user_id': keycloak_user_id,
}
if max_budget is not None:
json_data['max_budget_in_team'] = max_budget
response = await client.post( response = await client.post(
f'{LITE_LLM_API_URL}/team/member_update', f'{LITE_LLM_API_URL}/team/member_update',
json={ json=json_data,
'team_id': team_id,
'user_id': keycloak_user_id,
'max_budget_in_team': max_budget,
},
) )
# Failed to update user in team - this is an unforseen error state... # Failed to update user in team - this is an unforseen error state...
if not response.is_success: if not response.is_success:
logger.error( logger.error(