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

fix linting

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/_types.py
View file

@ -656,8 +656,12 @@ class UpdateKeyRequest(GenerateKeyRequest):
metadata: Optional[dict] = None 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 key: Optional[str] = None
duration: Optional[str] = None
spend: Optional[float] = None
metadata: Optional[dict] = None
class KeyRequest(LiteLLMBase): class KeyRequest(LiteLLMBase):

6
litellm/proxy/management_endpoints/key_management_endpoints.py
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) data_json: dict = data.dict(exclude_unset=True)
key = data_json.pop("key", None) key = data_json.pop("key", None)
@ -1055,7 +1057,7 @@ async def regenerate_key_fn(
# Update the token in the database # Update the token in the database
updated_token = await prisma_client.db.litellm_verificationtoken.update( updated_token = await prisma_client.db.litellm_verificationtoken.update(
where={"token": hashed_api_key}, where={"token": hashed_api_key},
data=update_data, data=update_data, # type: ignore
) )
updated_token_dict = {} updated_token_dict = {}