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

fix(proxy_server.py): enforce budget limit if global proxy limit reached

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-24 17:11:40 -08:00
parent 624da17698
commit 30a8071bf1
2 changed files with 70 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

35
litellm/proxy/utils.py
View file

@ -361,6 +361,7 @@ class PrismaClient:
self,
token: Optional[str] = None,
user_id: Optional[str] = None,
user_id_list: Optional[list] = None,
key_val: Optional[dict] = None,
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
query_type: Literal["find_unique", "find_all"] = "find_unique",
@ -442,6 +443,17 @@ class PrismaClient:
"budget_reset_at": {"lt": reset_at},
}
)
elif query_type == "find_all" and user_id_list is not None:
user_id_values = str(tuple(user_id_list))
sql_query = f"""
SELECT *
FROM "LiteLLM_UserTable"
WHERE "user_id" IN {user_id_values}
"""
# Execute the raw query
# The asterisk before `user_id_list` unpacks the list into separate arguments
response = await self.db.query_raw(sql_query)
return response
elif table_name == "user" and query_type == "find_all":
response = await self.db.litellm_usertable.find_many( # type: ignore
@ -586,6 +598,7 @@ class PrismaClient:
user_id: Optional[str] = None,
query_type: Literal["update", "update_many"] = "update",
table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
update_key_values: Optional[dict] = None,
):
"""
Update existing data
@ -612,28 +625,22 @@ class PrismaClient:
user_id is not None
or (table_name is not None and table_name == "user")
and query_type == "update"
and update_key_values is not None
):
"""
If data['spend'] + data['user'], update the user table with spend info as well
"""
if user_id is None:
user_id = db_data["user_id"]
update_user_row = await self.db.litellm_usertable.update(
update_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": user_id}, # type: ignore
data={**db_data}, # type: ignore
data={
"create": {**db_data}, # type: ignore
"update": {
**update_key_values # type: ignore
}, # just update user-specified values, if it already exists
},
)
if update_user_row is None:
# if the provided user does not exist, STILL Track this!
# make a new user with {"user_id": user_id, "spend": data['spend']}
db_data["user_id"] = user_id
update_user_row = await self.db.litellm_usertable.upsert(
where={"user_id": user_id}, # type: ignore
data={
"create": {**db_data}, # type: ignore
"update": {}, # don't do anything if it already exists
},
)
verbose_proxy_logger.info(
"\033[91m"
+ f"DB User Table - update succeeded {update_user_row}"