fix linting

This commit is contained in:
Ishaan Jaff 2024-09-06 16:54:43 -07:00
parent c5aa422847
commit 7f461dbf68
2 changed files with 9 additions and 3 deletions

View file

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

View file

@ -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 = {}