From d0953a13499dc69f4ea27c85e2707ce2b5a98eee Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 17:04:25 -0800 Subject: [PATCH] (fix) chore --- litellm/proxy/proxy_server.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ae3840cf1..17a213f513 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9,10 +9,6 @@ import warnings import importlib import warnings -import logging - -logging.getLogger("prisma").setLevel(logging.DEBUG) - def showwarning(message, category, filename, lineno, file=None, line=None): traceback_info = f"{filename}:{lineno}: {category.__name__}: {message}\n" @@ -1142,7 +1138,6 @@ async def update_database( ) # set cooldown on alert soft_budget_cooldown = True - # track cost per model, for the given key spend_per_model = existing_spend_obj.model_spend or {} current_model = kwargs.get("model") @@ -1158,7 +1153,11 @@ async def update_database( # Update the cost column for the given token await prisma_client.update_data( token=token, - data={"spend": new_spend, "model_spend": spend_per_model}, + data={ + "spend": new_spend, + "model_spend": spend_per_model, + "soft_budget_cooldown": soft_budget_cooldown, + }, ) valid_token = user_api_key_cache.get_cache(key=token) @@ -1212,7 +1211,6 @@ async def update_database( payload["spend"] = response_cost if prisma_client is not None: await prisma_client.insert_data(data=payload, table_name="spend") - elif custom_db_client is not None: await custom_db_client.insert_data(payload, table_name="spend") except Exception as e: