phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test(test_key_generate_prisma.py): fix check for raising error if key generate max budget > key upperbound

Browse source
This commit is contained in:
Krrish Dholakia 2024-04-09 13:07:34 -07:00
parent 23ebb5cd1d
commit 4495f6c058
2 changed files with 7 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -44,6 +44,7 @@ class LiteLLM_UpperboundKeyGenerateParams(LiteLLMBase):
""" """
max_budget: Optional[float] = None max_budget: Optional[float] = None
budget_duration: Optional[str] = None
max_parallel_requests: Optional[int] = None max_parallel_requests: Optional[int] = None
tpm_limit: Optional[int] = None tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None rpm_limit: Optional[int] = None

21
litellm/tests/test_key_generate_prisma.py
View file

@ -66,6 +66,7 @@ from litellm.proxy._types import (
GenerateKeyRequest, GenerateKeyRequest,
NewTeamRequest, NewTeamRequest,
UserAPIKeyAuth, UserAPIKeyAuth,
LiteLLM_UpperboundKeyGenerateParams,
) )
from litellm.proxy.utils import DBClient from litellm.proxy.utils import DBClient
from starlette.datastructures import URL from starlette.datastructures import URL
@ -1627,10 +1628,9 @@ async def test_upperbound_key_params(prisma_client):
""" """
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm.upperbound_key_generate_params = { litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
"max_budget": 0.001, max_budget=0.001, budget_duration="1m"
"budget_duration": "1m", )
}
await litellm.proxy.proxy_server.prisma_client.connect() await litellm.proxy.proxy_server.prisma_client.connect()
try: try:
request = GenerateKeyRequest( request = GenerateKeyRequest(
@ -1638,18 +1638,9 @@ async def test_upperbound_key_params(prisma_client):
budget_duration="30d", budget_duration="30d",
) )
key = await generate_key_fn(request) key = await generate_key_fn(request)
generated_key = key.key # print(result)
result = await info_key_fn(key=generated_key)
key_info = result["info"]
# assert it used the upper bound for max_budget, and budget_duration
assert key_info["max_budget"] == 0.001
assert key_info["budget_duration"] == "1m"
print(result)
except Exception as e: except Exception as e:
print("Got Exception", e) assert e.code == 400
pytest.fail(f"Got exception {e}")
def test_get_bearer_token(): def test_get_bearer_token():