phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat - ResetTeamBudgetRequest

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-25 16:52:41 -07:00
parent 920dc1754b
commit 7022993531
1 changed files with 66 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

69
litellm/proxy/utils.py
View file

@ -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,