diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0746b4e10..7c522f779 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d3c95f350..b14727d1a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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"