fix(dynamo_db.py): fix update bug

This commit is contained in:
Krrish Dholakia 2024-01-24 21:29:56 -08:00
parent 8e1157fc92
commit 0752048b81
3 changed files with 6 additions and 1 deletions

View file

@ -142,6 +142,7 @@ class GenerateRequestBase(LiteLLMBase):
class GenerateKeyRequest(GenerateRequestBase): class GenerateKeyRequest(GenerateRequestBase):
duration: Optional[str] = "1h" duration: Optional[str] = "1h"
aliases: Optional[dict] = {} aliases: Optional[dict] = {}
config: Optional[dict] = {}
class UpdateKeyRequest(GenerateKeyRequest): class UpdateKeyRequest(GenerateKeyRequest):

View file

@ -243,6 +243,10 @@ class DynamoDBWrapper(CustomDB):
and isinstance(v, str) and isinstance(v, str)
): ):
new_response[k] = json.loads(v) new_response[k] = json.loads(v)
elif (k == "tpm_limit" or k == "rpm_limit") and isinstance(
v, float
):
new_response[k] = int(v)
else: else:
new_response[k] = v new_response[k] = v
new_response = LiteLLM_VerificationToken(**new_response) new_response = LiteLLM_VerificationToken(**new_response)

View file

@ -686,7 +686,7 @@ async def update_database(
continue continue
if prisma_client is not None: if prisma_client is not None:
existing_spend_obj = await prisma_client.get_data(user_id=id) existing_spend_obj = await prisma_client.get_data(user_id=id)
elif custom_db_client is not None: elif custom_db_client is not None and id != litellm_proxy_budget_name:
existing_spend_obj = await custom_db_client.get_data( existing_spend_obj = await custom_db_client.get_data(
key=id, table_name="user" key=id, table_name="user"
) )