diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b710165cb44..70aba558094 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -14,6 +14,7 @@ from litellm.proxy._types import ( CallInfo, WebhookEvent, AlertType, + ResetTeamBudgetRequest, ) from litellm.caching import DualCache, RedisCache from litellm.router import Deployment, ModelInfo, LiteLLM_Params @@ -1021,7 +1022,7 @@ class PrismaClient: return response elif table_name == "spend": verbose_proxy_logger.debug( - f"PrismaClient: get_data: table_name == 'spend'" + "PrismaClient: get_data: table_name == 'spend'" ) if key_val is not None: if query_type == "find_unique": @@ -1047,6 +1048,12 @@ class PrismaClient: response = await self.db.litellm_teamtable.find_unique( where={"team_id": team_id} # type: ignore ) + elif query_type == "find_all" and reset_at is not None: + response = await self.db.litellm_teamtable.find_many( + where={ # type:ignore + "budget_reset_at": {"lt": reset_at}, + } + ) elif query_type == "find_all" and user_id is not None: response = await self.db.litellm_teamtable.find_many( where={ @@ -1438,7 +1445,7 @@ class PrismaClient: ) await batcher.commit() print_verbose( - "\033[91m" + f"DB Token Table update succeeded" + "\033[0m" + "\033[91m" + "DB Token Table update succeeded" + "\033[0m" ) elif ( table_name is not None @@ -1467,8 +1474,40 @@ class PrismaClient: ) await batcher.commit() verbose_proxy_logger.info( - "\033[91m" + f"DB User Table Batch update succeeded" + "\033[0m" + "\033[91m" + "DB User Table Batch update succeeded" + "\033[0m" ) + elif ( + table_name is not None + and table_name == "team" + and query_type == "update_many" + and data_list is not None + and isinstance(data_list, list) + ): + # Batch write update queries + batcher = self.db.batch_() + for idx, team in enumerate(data_list): + try: + data_json = self.jsonify_object( + data=team.model_dump(exclude_none=True) + ) + except: + data_json = self.jsonify_object( + data=team.dict(exclude_none=True) + ) + batcher.litellm_teamtable.upsert( + where={"team_id": team.team_id}, # type: ignore + data={ + "create": {**data_json}, # type: ignore + "update": { + **data_json # type: ignore + }, # just update user-specified values, if it already exists + }, + ) + await batcher.commit() + verbose_proxy_logger.info( + "\033[91m" + "DB Team Table Batch update succeeded" + "\033[0m" + ) + except Exception as e: import traceback @@ -2043,6 +2082,30 @@ async def reset_budget(prisma_client: PrismaClient): query_type="update_many", data_list=users_to_reset, table_name="user" ) + ## Reset Team Budget + now = datetime.utcnow() + teams_to_reset = await prisma_client.get_data( + table_name="team", + query_type="find_all", + reset_at=now, + ) + + if teams_to_reset is not None and len(teams_to_reset) > 0: + team_reset_requests = [] + for team in teams_to_reset: + duration_s = _duration_in_seconds(duration=team.budget_duration) + reset_team_budget_request = ResetTeamBudgetRequest( + team_id=team.team_id, + spend=0.0, + budget_reset_at=now + timedelta(seconds=duration_s), + ) + team_reset_requests.append(reset_team_budget_request) + await prisma_client.update_data( + query_type="update_many", + data_list=team_reset_requests, + table_name="team", + ) + async def update_spend( prisma_client: PrismaClient,