phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

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

Browse source
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
Download patch file Download diff file Expand all files Collapse all files

28
litellm/proxy/proxy_server.py
View file

@ -1233,10 +1233,11 @@ async def update_database(
user_ids.append(litellm_proxy_budget_name)
### KEY CHANGE ###
for _id in user_ids:
prisma_client.user_list_transactons[_id] = (
response_cost
+ prisma_client.user_list_transactons.get(_id, 0)
)
if _id is not None:
prisma_client.user_list_transactons[_id] = (
response_cost
+ prisma_client.user_list_transactons.get(_id, 0)
)
if end_user_id is not None:
prisma_client.end_user_list_transactons[end_user_id] = (
response_cost
@ -1364,16 +1365,17 @@ async def update_database(
)
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)
# if db_writer_client is not None:
# print("Tries to make call")
# 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}")
elif prisma_client is not None:
await prisma_client.insert_data(data=payload, table_name="spend")
except Exception as e:
verbose_proxy_logger.debug(
f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}"

24
litellm/proxy/utils.py
View file

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