From 0752048b8147c94b2bb6e1edcb03e144aa1f0ee7 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 24 Jan 2024 21:29:56 -0800 Subject: [PATCH] fix(dynamo_db.py): fix update bug --- litellm/proxy/_types.py | 1 + litellm/proxy/db/dynamo_db.py | 4 ++++ litellm/proxy/proxy_server.py | 2 +- 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e64fc729e95..13c0862853b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -142,6 +142,7 @@ class GenerateRequestBase(LiteLLMBase): class GenerateKeyRequest(GenerateRequestBase): duration: Optional[str] = "1h" aliases: Optional[dict] = {} + config: Optional[dict] = {} class UpdateKeyRequest(GenerateKeyRequest): diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 83cf6b15724..534adbddc94 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -243,6 +243,10 @@ class DynamoDBWrapper(CustomDB): and isinstance(v, str) ): new_response[k] = json.loads(v) + elif (k == "tpm_limit" or k == "rpm_limit") and isinstance( + v, float + ): + new_response[k] = int(v) else: new_response[k] = v new_response = LiteLLM_VerificationToken(**new_response) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 027e29f94d4..c13849e2a32 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -686,7 +686,7 @@ async def update_database( continue if prisma_client is not None: 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( key=id, table_name="user" )