forked from phoenix/litellm-mirror
(feat) track spend per team
This commit is contained in:
ishaan-jaff
parent
039617766f
commit
88037b68b0
2 changed files with 95 additions and 2 deletions
|
|
@ -843,6 +843,7 @@ async def _PROXY_track_cost_callback(
|
|||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
kwargs=kwargs,
|
||||
completion_response=completion_response,
|
||||
start_time=start_time,
|
||||
|
|
@ -880,6 +881,7 @@ async def update_database(
|
|||
token,
|
||||
response_cost,
|
||||
user_id=None,
|
||||
team_id=None,
|
||||
kwargs=None,
|
||||
completion_response=None,
|
||||
start_time=None,
|
||||
|
|
@ -887,7 +889,7 @@ async def update_database(
|
|||
):
|
||||
try:
|
||||
verbose_proxy_logger.info(
|
||||
f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}"
|
||||
f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}"
|
||||
)
|
||||
|
||||
### [TODO] STEP 1: GET KEY + USER SPEND ### (key, user)
|
||||
|
|
@ -1037,8 +1039,69 @@ async def update_database(
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.info(f"Update Spend Logs DB failed to execute")
|
||||
|
||||
### UPDATE KEY SPEND ###
|
||||
async def _update_team_db():
|
||||
try:
|
||||
verbose_proxy_logger.debug(
|
||||
f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}."
|
||||
)
|
||||
if team_id is None:
|
||||
verbose_proxy_logger.debug(
|
||||
"track_cost_callback: team_id is None. Not tracking spend for team"
|
||||
)
|
||||
return
|
||||
if prisma_client is not None:
|
||||
# Fetch the existing cost for the given token
|
||||
existing_spend_obj = await prisma_client.get_data(
|
||||
team_id=team_id, table_name="team"
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_team_db: existing spend: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
verbose_proxy_logger.debug(f"new cost: {new_spend}")
|
||||
# Update the cost column for the given token
|
||||
await prisma_client.update_data(
|
||||
team_id=team_id, data={"spend": new_spend}, table_name="team"
|
||||
)
|
||||
|
||||
elif custom_db_client is not None:
|
||||
# Fetch the existing cost for the given token
|
||||
existing_spend_obj = await custom_db_client.get_data(
|
||||
key=token, table_name="key"
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
f"_update_key_db existing spend: {existing_spend_obj}"
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
existing_spend = 0
|
||||
else:
|
||||
existing_spend = existing_spend_obj.spend
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
new_spend = existing_spend + response_cost
|
||||
|
||||
verbose_proxy_logger.debug(f"new cost: {new_spend}")
|
||||
# Update the cost column for the given token
|
||||
await custom_db_client.update_data(
|
||||
key=token, value={"spend": new_spend}, table_name="key"
|
||||
)
|
||||
|
||||
valid_token = user_api_key_cache.get_cache(key=token)
|
||||
if valid_token is not None:
|
||||
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")
|
||||
|
||||
asyncio.create_task(_update_user_db())
|
||||
asyncio.create_task(_update_key_db())
|
||||
asyncio.create_task(_update_team_db())
|
||||
asyncio.create_task(_insert_spend_log_to_db())
|
||||
verbose_proxy_logger.info("Successfully updated spend in all 3 tables")
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -808,8 +808,9 @@ class PrismaClient:
|
|||
data: dict = {},
|
||||
data_list: Optional[List] = None,
|
||||
user_id: Optional[str] = None,
|
||||
team_id: Optional[str] = None,
|
||||
query_type: Literal["update", "update_many"] = "update",
|
||||
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
|
||||
table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None,
|
||||
update_key_values: Optional[dict] = None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -860,6 +861,35 @@ class PrismaClient:
|
|||
+ "\033[0m"
|
||||
)
|
||||
return {"user_id": user_id, "data": db_data}
|
||||
elif (
|
||||
team_id is not None
|
||||
or (table_name is not None and table_name == "team")
|
||||
and query_type == "update"
|
||||
):
|
||||
"""
|
||||
If data['spend'] + data['user'], update the user table with spend info as well
|
||||
"""
|
||||
if team_id is None:
|
||||
team_id = db_data["team_id"]
|
||||
if update_key_values is None:
|
||||
update_key_values = db_data
|
||||
if "team_id" not in db_data and team_id is not None:
|
||||
db_data["team_id"] = team_id
|
||||
update_team_row = await self.db.litellm_teamtable.upsert(
|
||||
where={"team_id": team_id}, # type: ignore
|
||||
data={
|
||||
"create": {**db_data}, # type: ignore
|
||||
"update": {
|
||||
**update_key_values # type: ignore
|
||||
}, # just update user-specified values, if it already exists
|
||||
},
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"\033[91m"
|
||||
+ f"DB Team Table - update succeeded {update_team_row}"
|
||||
+ "\033[0m"
|
||||
)
|
||||
return {"team_id": team_id, "data": db_data}
|
||||
elif (
|
||||
table_name is not None
|
||||
and table_name == "key"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue