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

(fix) fix batch update user db

Browse source
This commit is contained in:
ishaan-jaff 2024-03-05 16:46:58 -08:00
parent 78d87a4fbd
commit 3f7bf5c6b1
1 changed files with 17 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/proxy_server.py
View file

@ -1067,6 +1067,11 @@ async def update_database(
) )
data_list.append(existing_spend_obj) data_list.append(existing_spend_obj)
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 # Update the cost column for the given user id
if prisma_client is not None: if prisma_client is not None:
await prisma_client.update_data( await prisma_client.update_data(
@ -1074,13 +1079,10 @@ async def update_database(
query_type="update_many", query_type="update_many",
table_name="user", table_name="user",
) )
elif 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"
)
except Exception as e: 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 ### ### UPDATE KEY SPEND ###
async def _update_key_db(): async def _update_key_db():
@ -1215,7 +1217,9 @@ async def update_database(
await custom_db_client.insert_data(payload, table_name="spend") await custom_db_client.insert_data(payload, table_name="spend")
except Exception as e: 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 ### ### UPDATE KEY SPEND ###
async def _update_team_db(): async def _update_team_db():
@ -1286,7 +1290,9 @@ async def update_database(
valid_token.spend = new_spend valid_token.spend = new_spend
user_api_key_cache.set_cache(key=token, value=valid_token) user_api_key_cache.set_cache(key=token, value=valid_token)
except Exception as e: 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_user_db())
asyncio.create_task(_update_key_db()) asyncio.create_task(_update_key_db())