phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Feat-Proxy] Add upperbound key duration param (#5727)

Browse source 
* add upperbound key duration param

* use upper bound values when None set

* docs upperbound params
This commit is contained in:
Ishaan Jaff 2024-09-16 16:28:36 -07:00 • committed by GitHub
parent 0d18292549
commit a252c95c7f
5 changed files with 147 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

50
litellm/tests/test_key_generate_prisma.py
View file

@ -2045,7 +2045,7 @@ async def test_default_key_params(prisma_client):
@pytest.mark.asyncio()
async def test_upperbound_key_params(prisma_client):
async def test_upperbound_key_param_larger_budget(prisma_client):
"""
- create key
- get key info
@ -2068,6 +2068,54 @@ async def test_upperbound_key_params(prisma_client):
assert e.code == str(400)
@pytest.mark.asyncio()
async def test_upperbound_key_param_larger_duration(prisma_client):
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
max_budget=100, duration="14d"
)
await litellm.proxy.proxy_server.prisma_client.connect()
try:
request = GenerateKeyRequest(
max_budget=10,
duration="30d",
)
key = await generate_key_fn(request)
pytest.fail("Expected this to fail but it passed")
# print(result)
except Exception as e:
assert e.code == str(400)
@pytest.mark.asyncio()
async def test_upperbound_key_param_none_duration(prisma_client):
from datetime import datetime, timedelta
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
max_budget=100, duration="14d"
)
await litellm.proxy.proxy_server.prisma_client.connect()
try:
request = GenerateKeyRequest()
key = await generate_key_fn(request)
print(key)
# print(result)
assert key.max_budget == 100
assert key.expires is not None
_date_key_expires = key.expires.date()
_fourteen_days_from_now = (datetime.now() + timedelta(days=14)).date()
assert _date_key_expires == _fourteen_days_from_now
except Exception as e:
pytest.fail(f"Got exception {e}")
def test_get_bearer_token():
from litellm.proxy.auth.user_api_key_auth import _get_bearer_token