From 3f7bf5c6b19e746a6609c48650d42cd74ea8ce04 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 16:46:58 -0800 Subject: [PATCH] (fix) fix batch update user db --- litellm/proxy/proxy_server.py | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 47b8c1535..628f55852 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1067,20 +1067,22 @@ async def update_database( ) data_list.append(existing_spend_obj) - # Update the cost column for the given user id - if prisma_client is not None: - await prisma_client.update_data( - data_list=data_list, - query_type="update_many", - table_name="user", - ) - elif custom_db_client is not None and user_id is not None: + if custom_db_client is not None and user_id is not None: new_spend = data_list[0].spend await custom_db_client.update_data( key=user_id, value={"spend": new_spend}, table_name="user" ) + # Update the cost column for the given user id + if prisma_client is not None: + await prisma_client.update_data( + data_list=data_list, + query_type="update_many", + table_name="user", + ) except Exception as e: - verbose_proxy_logger.info(f"Update User DB call failed to execute") + verbose_proxy_logger.info( + f"Update User DB call failed to execute {str(e)}" + ) ### UPDATE KEY SPEND ### async def _update_key_db(): @@ -1215,7 +1217,9 @@ async def update_database( await custom_db_client.insert_data(payload, table_name="spend") except Exception as e: - verbose_proxy_logger.info(f"Update Spend Logs DB failed to execute") + verbose_proxy_logger.info( + f"Update Spend Logs DB failed to execute - {str(e)}" + ) ### UPDATE KEY SPEND ### async def _update_team_db(): @@ -1286,7 +1290,9 @@ async def update_database( valid_token.spend = new_spend user_api_key_cache.set_cache(key=token, value=valid_token) except Exception as e: - verbose_proxy_logger.info(f"Update Team DB failed to execute") + verbose_proxy_logger.info( + f"Update Team DB failed to execute - {str(e)}" + ) asyncio.create_task(_update_user_db()) asyncio.create_task(_update_key_db())