fix: use more descriptive flag

This commit is contained in:
Krrish Dholakia 2024-08-07 18:59:46 -07:00
parent b7e31638fd
commit ec0b511119
3 changed files with 5 additions and 4 deletions

View file

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

View file

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

View file

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