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

use simple static methods for updating spend

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-27 16:57:12 -07:00
parent 7995fd7c98
commit 403f2ef68d
1 changed files with 201 additions and 148 deletions
Download patch file Download diff file Expand all files Collapse all files

349
litellm/proxy/db/db_spend_update_writer.py
View file

@ -13,6 +13,7 @@ from typing import Any, Optional, Union
import litellm import litellm
from litellm._logging import verbose_proxy_logger from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token
@ -27,7 +28,7 @@ class DBSpendUpdateWriter:
""" """
@staticmethod @staticmethod
async def update_database( # noqa: PLR0915 async def update_database(
# LiteLLM management object fields # LiteLLM management object fields
token: Optional[str], token: Optional[str],
user_id: Optional[str], user_id: Optional[str],
@ -59,155 +60,47 @@ class DBSpendUpdateWriter:
else: else:
hashed_token = token hashed_token = token
### UPDATE USER SPEND ### asyncio.create_task(
async def _update_user_db(): DBSpendUpdateWriter._update_user_db(
""" response_cost=response_cost,
- Update that user's row user_id=user_id,
- Update litellm-proxy-budget row (global proxy spend) prisma_client=prisma_client,
""" user_api_key_cache=user_api_key_cache,
## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db litellm_proxy_budget_name=litellm_proxy_budget_name,
existing_user_obj = await user_api_key_cache.async_get_cache( end_user_id=end_user_id,
key=user_id
) )
if existing_user_obj is not None and isinstance( )
existing_user_obj, dict asyncio.create_task(
): DBSpendUpdateWriter._update_key_db(
existing_user_obj = LiteLLM_UserTable(**existing_user_obj) response_cost=response_cost,
try: hashed_token=hashed_token,
if prisma_client is not None: # update prisma_client=prisma_client,
user_ids = [user_id] )
if ( )
litellm.max_budget > 0 asyncio.create_task(
): # track global proxy budget, if user set max budget DBSpendUpdateWriter._update_team_db(
user_ids.append(litellm_proxy_budget_name) response_cost=response_cost,
### KEY CHANGE ### team_id=team_id,
for _id in user_ids: user_id=user_id,
if _id is not None: prisma_client=prisma_client,
prisma_client.user_list_transactons[_id] = ( )
response_cost )
+ prisma_client.user_list_transactons.get(_id, 0) asyncio.create_task(
) DBSpendUpdateWriter._update_org_db(
if end_user_id is not None: response_cost=response_cost,
prisma_client.end_user_list_transactons[end_user_id] = ( org_id=org_id,
response_cost prisma_client=prisma_client,
+ prisma_client.end_user_list_transactons.get( )
end_user_id, 0 )
)
)
except Exception as e:
verbose_proxy_logger.info(
"\033[91m"
+ f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}"
)
### UPDATE KEY SPEND ###
async def _update_key_db():
try:
verbose_proxy_logger.debug(
f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}."
)
if hashed_token is None:
return
if prisma_client is not None:
prisma_client.key_list_transactons[hashed_token] = (
response_cost
+ prisma_client.key_list_transactons.get(hashed_token, 0)
)
except Exception as e:
verbose_proxy_logger.exception(
f"Update Key DB Call failed to execute - {str(e)}"
)
raise e
### UPDATE SPEND LOGS ###
async def _insert_spend_log_to_db():
try:
if prisma_client:
payload = get_logging_payload(
kwargs=kwargs,
response_obj=completion_response,
start_time=start_time,
end_time=end_time,
)
payload["spend"] = response_cost or 0.0
DBSpendUpdateWriter._set_spend_logs_payload(
payload=payload,
spend_logs_url=os.getenv("SPEND_LOGS_URL"),
prisma_client=prisma_client,
)
except Exception as e:
verbose_proxy_logger.debug(
f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}"
)
raise e
### UPDATE TEAM SPEND ###
async def _update_team_db():
try:
verbose_proxy_logger.debug(
f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}."
)
if team_id is None:
verbose_proxy_logger.debug(
"track_cost_callback: team_id is None. Not tracking spend for team"
)
return
if prisma_client is not None:
prisma_client.team_list_transactons[team_id] = (
response_cost
+ prisma_client.team_list_transactons.get(team_id, 0)
)
try:
# Track spend of the team member within this team
# key is "team_id::<value>::user_id::<value>"
team_member_key = f"team_id::{team_id}::user_id::{user_id}"
prisma_client.team_member_list_transactons[
team_member_key
] = (
response_cost
+ prisma_client.team_member_list_transactons.get(
team_member_key, 0
)
)
except Exception:
pass
except Exception as e:
verbose_proxy_logger.info(
f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}"
)
raise e
### UPDATE ORG SPEND ###
async def _update_org_db():
try:
verbose_proxy_logger.debug(
"adding spend to org db. Response cost: {}. org_id: {}.".format(
response_cost, org_id
)
)
if org_id is None:
verbose_proxy_logger.debug(
"track_cost_callback: org_id is None. Not tracking spend for org"
)
return
if prisma_client is not None:
prisma_client.org_list_transactons[org_id] = (
response_cost
+ prisma_client.org_list_transactons.get(org_id, 0)
)
except Exception as e:
verbose_proxy_logger.info(
f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}"
)
raise e
asyncio.create_task(_update_user_db())
asyncio.create_task(_update_key_db())
asyncio.create_task(_update_team_db())
asyncio.create_task(_update_org_db())
if disable_spend_logs is False: if disable_spend_logs is False:
await _insert_spend_log_to_db() await DBSpendUpdateWriter._insert_spend_log_to_db(
kwargs=kwargs,
completion_response=completion_response,
start_time=start_time,
end_time=end_time,
response_cost=response_cost,
prisma_client=prisma_client,
)
else: else:
verbose_proxy_logger.info( verbose_proxy_logger.info(
"disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur."
@ -219,6 +112,166 @@ class DBSpendUpdateWriter:
f"Error updating Prisma database: {traceback.format_exc()}" f"Error updating Prisma database: {traceback.format_exc()}"
) )
@staticmethod
async def _update_key_db(
response_cost: Optional[float],
hashed_token: Optional[str],
prisma_client: Optional[PrismaClient],
):
try:
verbose_proxy_logger.debug(
f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}."
)
if hashed_token is None:
return
if prisma_client is not None:
prisma_client.key_list_transactons[hashed_token] = (
response_cost
+ prisma_client.key_list_transactons.get(hashed_token, 0)
)
except Exception as e:
verbose_proxy_logger.exception(
f"Update Key DB Call failed to execute - {str(e)}"
)
raise e
@staticmethod
async def _update_user_db(
response_cost: Optional[float],
user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
litellm_proxy_budget_name: Optional[str],
end_user_id: Optional[str] = None,
):
"""
- Update that user's row
- Update litellm-proxy-budget row (global proxy spend)
"""
## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db
existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
if existing_user_obj is not None and isinstance(existing_user_obj, dict):
existing_user_obj = LiteLLM_UserTable(**existing_user_obj)
try:
if prisma_client is not None: # update
user_ids = [user_id]
if (
litellm.max_budget > 0
): # track global proxy budget, if user set max budget
user_ids.append(litellm_proxy_budget_name)
### KEY CHANGE ###
for _id in user_ids:
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
+ prisma_client.end_user_list_transactons.get(end_user_id, 0)
)
except Exception as e:
verbose_proxy_logger.info(
"\033[91m"
+ f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}"
)
@staticmethod
async def _update_team_db(
response_cost: Optional[float],
team_id: Optional[str],
user_id: Optional[str],
prisma_client: Optional[PrismaClient],
):
try:
verbose_proxy_logger.debug(
f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}."
)
if team_id is None:
verbose_proxy_logger.debug(
"track_cost_callback: team_id is None. Not tracking spend for team"
)
return
if prisma_client is not None:
prisma_client.team_list_transactons[team_id] = (
response_cost + prisma_client.team_list_transactons.get(team_id, 0)
)
try:
# Track spend of the team member within this team
# key is "team_id::<value>::user_id::<value>"
team_member_key = f"team_id::{team_id}::user_id::{user_id}"
prisma_client.team_member_list_transactons[team_member_key] = (
response_cost
+ prisma_client.team_member_list_transactons.get(
team_member_key, 0
)
)
except Exception:
pass
except Exception as e:
verbose_proxy_logger.info(
f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}"
)
raise e
@staticmethod
async def _update_org_db(
response_cost: Optional[float],
org_id: Optional[str],
prisma_client: Optional[PrismaClient],
):
try:
verbose_proxy_logger.debug(
"adding spend to org db. Response cost: {}. org_id: {}.".format(
response_cost, org_id
)
)
if org_id is None:
verbose_proxy_logger.debug(
"track_cost_callback: org_id is None. Not tracking spend for org"
)
return
if prisma_client is not None:
prisma_client.org_list_transactons[org_id] = (
response_cost + prisma_client.org_list_transactons.get(org_id, 0)
)
except Exception as e:
verbose_proxy_logger.info(
f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}"
)
raise e
@staticmethod
async def _insert_spend_log_to_db(
kwargs: Optional[dict],
completion_response: Optional[Union[litellm.ModelResponse, Any, Exception]],
start_time: Optional[datetime],
end_time: Optional[datetime],
response_cost: Optional[float],
prisma_client: Optional[PrismaClient],
):
try:
if prisma_client:
payload = get_logging_payload(
kwargs=kwargs,
response_obj=completion_response,
start_time=start_time,
end_time=end_time,
)
payload["spend"] = response_cost or 0.0
DBSpendUpdateWriter._set_spend_logs_payload(
payload=payload,
spend_logs_url=os.getenv("SPEND_LOGS_URL"),
prisma_client=prisma_client,
)
except Exception as e:
verbose_proxy_logger.debug(
f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}"
)
raise e
@staticmethod @staticmethod
def _set_spend_logs_payload( def _set_spend_logs_payload(
payload: Union[dict, SpendLogsPayload], payload: Union[dict, SpendLogsPayload],