diff --git a/litellm/__init__.py b/litellm/__init__.py index dfc3f3fc1b..b284525d74 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -259,6 +259,7 @@ upperbound_key_generate_params: Optional[LiteLLM_UpperboundKeyGenerateParams] = default_user_params: Optional[Dict] = None default_team_settings: Optional[List] = None max_user_budget: Optional[float] = None +max_internal_user_budget: Optional[float] = None max_end_user_budget: Optional[float] = None #### REQUEST PRIORITIZATION #### priority_reservation: Optional[Dict[str, float]] = None diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index de109db46b..f54eca9a2f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -88,8 +88,8 @@ async def new_user( ) if "max_budget" in data_json and data_json["max_budget"] is None: - if litellm.max_user_budget is not None: - data_json["max_budget"] = litellm.max_user_budget + if litellm.max_internal_user_budget is not None: + data_json["max_budget"] = litellm.max_internal_user_budget response = await generate_key_helper_fn(request_type="user", **data_json) diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index 35198af807..a7bc1c4a22 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -816,7 +816,7 @@ async def test_create_user_default_budget(prisma_client): setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - setattr(litellm, "max_user_budget", 10) + setattr(litellm, "max_internal_user_budget", 10) await litellm.proxy.proxy_server.prisma_client.connect() user = f"ishaan {uuid.uuid4().hex}" request = NewUserRequest(user_id=user) # create a key with no budget @@ -834,5 +834,5 @@ async def test_create_user_default_budget(prisma_client): assert ( mock_client.call_args.kwargs["data"]["max_budget"] - == litellm.max_user_budget + == litellm.max_internal_user_budget )