From 7f461dbf68f3e6551236a389f1aa5f70c1c9fb61 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 6 Sep 2024 16:54:43 -0700 Subject: [PATCH] fix linting --- litellm/proxy/_types.py | 6 +++++- .../proxy/management_endpoints/key_management_endpoints.py | 6 ++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 812b572fb..2bbd14a16 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -656,8 +656,12 @@ class UpdateKeyRequest(GenerateKeyRequest): metadata: Optional[dict] = None -class RegenerateKeyRequest(UpdateKeyRequest): +class RegenerateKeyRequest(GenerateKeyRequest): + # This needs to be different from UpdateKeyRequest, because "key" is optional for this key: Optional[str] = None + duration: Optional[str] = None + spend: Optional[float] = None + metadata: Optional[dict] = None class KeyRequest(LiteLLMBase): diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 43f4c1fcb..553cdb177 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -280,7 +280,9 @@ async def generate_key_fn( ) -async def prepare_key_update_data(data: UpdateKeyRequest, existing_key_row): +async def prepare_key_update_data( + data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row +): data_json: dict = data.dict(exclude_unset=True) key = data_json.pop("key", None) @@ -1055,7 +1057,7 @@ async def regenerate_key_fn( # Update the token in the database updated_token = await prisma_client.db.litellm_verificationtoken.update( where={"token": hashed_api_key}, - data=update_data, + data=update_data, # type: ignore ) updated_token_dict = {}