fix(proxy_server.py): allow user to pass in spend logs collector url

This commit is contained in:
Krrish Dholakia 2024-03-28 09:14:30 -07:00
parent 1e856443e1
commit 082f1e4085
2 changed files with 26 additions and 26 deletions

View file

@ -1233,10 +1233,11 @@ async def update_database(
user_ids.append(litellm_proxy_budget_name) user_ids.append(litellm_proxy_budget_name)
### KEY CHANGE ### ### KEY CHANGE ###
for _id in user_ids: for _id in user_ids:
prisma_client.user_list_transactons[_id] = ( if _id is not None:
response_cost prisma_client.user_list_transactons[_id] = (
+ prisma_client.user_list_transactons.get(_id, 0) response_cost
) + prisma_client.user_list_transactons.get(_id, 0)
)
if end_user_id is not None: if end_user_id is not None:
prisma_client.end_user_list_transactons[end_user_id] = ( prisma_client.end_user_list_transactons[end_user_id] = (
response_cost response_cost
@ -1364,16 +1365,17 @@ async def update_database(
) )
payload["spend"] = response_cost payload["spend"] = response_cost
if prisma_client is not None: if (
os.getenv("SPEND_LOGS_URL", None) is not None
and prisma_client is not None
):
if isinstance(payload["startTime"], datetime):
payload["startTime"] = payload["startTime"].isoformat()
if isinstance(payload["endTime"], datetime):
payload["endTime"] = payload["endTime"].isoformat()
prisma_client.spend_log_transactons.append(payload) prisma_client.spend_log_transactons.append(payload)
# if db_writer_client is not None: elif prisma_client is not None:
# print("Tries to make call") await prisma_client.insert_data(data=payload, table_name="spend")
# response = await db_writer_client.post(
# url="http://0.0.0.0:3000/spend/update",
# data=json.dumps(payload),
# headers={"Content-Type": "application/json"},
# )
# print(f"response: {response}")
except Exception as e: except Exception as e:
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}" f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}"

View file

@ -1717,11 +1717,6 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time):
# hash the api_key # hash the api_key
api_key = hash_token(api_key) api_key = hash_token(api_key)
# jsonify datetime object
# if isinstance(start_time, datetime):
# start_time = start_time.isoformat()
# if isinstance(end_time, datetime):
# end_time = end_time.isoformat()
# clean up litellm metadata # clean up litellm metadata
if isinstance(metadata, dict): if isinstance(metadata, dict):
clean_metadata = {} clean_metadata = {}
@ -2000,14 +1995,17 @@ async def update_spend(prisma_client: PrismaClient, db_writer_client: HTTPHandle
raise e raise e
### UPDATE SPEND LOGS ### ### UPDATE SPEND LOGS ###
# if len(prisma_client.spend_log_transactons) > 0: base_url = os.getenv("SPEND_LOGS_URL", None)
# response = await db_writer_client.post( if len(prisma_client.spend_log_transactons) > 0 and base_url is not None:
# url="http://0.0.0.0:3000/spend/update", if not base_url.endswith("/"):
# data=prisma_client.spend_log_transactons, base_url += "/"
# headers={"Content-Type": "application/json"}, response = await db_writer_client.post(
# ) url=base_url + "spend/update",
# if response.status_code == 200: data=json.dumps(prisma_client.spend_log_transactons), # type: ignore
# prisma_client.spend_log_transactons = [] headers={"Content-Type": "application/json"},
)
if response.status_code == 200:
prisma_client.spend_log_transactons = []
# async def monitor_spend_list(prisma_client: PrismaClient): # async def monitor_spend_list(prisma_client: PrismaClient):