Merge branch 'main' into litellm_use_helper_to_get_httpx_clients

This commit is contained in:
Ishaan Jaff 2024-09-10 15:02:54 -07:00 committed by GitHub
commit 3ebff903c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 58 additions and 4 deletions

View file

@ -293,9 +293,9 @@ async def prepare_key_update_data(
continue
if v is not None and v not in ([], {}, 0):
non_default_values[k] = v
if "duration" in non_default_values:
duration = non_default_values.pop("duration")
if duration and (isinstance(duration, str)) and len(duration) > 0:
duration_s = _duration_in_seconds(duration=duration)
expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
non_default_values["expires"] = expires

View file

@ -269,3 +269,56 @@ async def test_regenerate_api_key_with_new_alias_and_expiration(prisma_client):
now = datetime.now(dt.timezone.utc)
assert new_key.expires > now + dt.timedelta(days=29)
assert new_key.expires < now + dt.timedelta(days=31)
@pytest.mark.asyncio()
async def test_regenerate_key_ui(prisma_client):
litellm.set_verbose = True
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
import uuid
# generate new key
key_alias = f"test_alias_regenerate_key-{uuid.uuid4()}"
spend = 100
max_budget = 400
models = ["fake-openai-endpoint"]
new_key = await generate_key_fn(
data=GenerateKeyRequest(
key_alias=key_alias, spend=spend, max_budget=max_budget, models=models
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
),
)
generated_key = new_key.key
print(generated_key)
# assert the new key works as expected
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
async def return_body():
return_string = f'{{"model": "fake-openai-endpoint"}}'
# return string as bytes
return return_string.encode()
request.body = return_body
result = await user_api_key_auth(request=request, api_key=f"Bearer {generated_key}")
print(result)
# regenerate the key
new_key = await regenerate_key_fn(
key=generated_key,
data=RegenerateKeyRequest(duration=""),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="1234",
),
)
print("response from regenerate_key_fn", new_key)

View file

@ -89,6 +89,7 @@ from litellm.caching import DualCache
from litellm.proxy._types import (
DynamoDBArgs,
GenerateKeyRequest,
RegenerateKeyRequest,
KeyRequest,
LiteLLM_UpperboundKeyGenerateParams,
NewCustomerRequest,