From ddd9ca86a79e5cd95a846386a08e2af81c3cd0d1 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 18 Jan 2024 17:44:39 -0800 Subject: [PATCH] (feat) proxy - track cost for user_ids that do not exist --- litellm/proxy/utils.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 23b66f22d7..9f183644da 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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(