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) proxy - track cost for user_ids that do not exist

Browse source
This commit is contained in:
ishaan-jaff 2024-01-18 17:44:39 -08:00
parent a8ba5df90e
commit ddd9ca86a7
1 changed files with 22 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/utils.py
View file

@ -530,7 +530,11 @@ class PrismaClient:
where={"token": token}, # type: ignore
data={**db_data}, # type: ignore
)
print_verbose("\033[91m" + f"DB write succeeded {response}" + "\033[0m")
print_verbose(
"\033[91m"
+ f"DB Token Table update succeeded {response}"
+ "\033[0m"
)
return {"token": token, "data": db_data}
elif user_id is not None:
"""
@ -540,6 +544,23 @@ class PrismaClient:
where={"user_id": user_id}, # type: ignore
data={**db_data}, # type: ignore
)
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
},
)
print_verbose(
"\033[91m"
+ f"DB User Table - update succeeded {update_user_row}"
+ "\033[0m"
)
return {"user_id": user_id, "data": db_data}
except Exception as e:
asyncio.create_task(