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

(feat) working budgets per key

Browse source
This commit is contained in:
ishaan-jaff 2024-01-22 12:13:19 -08:00
parent de1502658d
commit 8979b74d49
1 changed files with 22 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/proxy_server.py
View file

@ -306,6 +306,7 @@ async def user_api_key_auth(
# 1. If token can call model # 1. If token can call model
# 2. If user_id for this token is in budget # 2. If user_id for this token is in budget
# 3. If token is expired # 3. If token is expired
# 4. If token spend is under Budget for the token
# Check 1. If token can call model # Check 1. If token can call model
litellm.model_alias_map = valid_token.aliases litellm.model_alias_map = valid_token.aliases
@ -406,6 +407,13 @@ async def user_api_key_auth(
detail=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", detail=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
) )
# Check 4. Token Spend is under budget
if valid_token.spend is not None and valid_token.max_budget is not None:
if valid_token.spend > valid_token.max_budget:
raise Exception(
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Token: {valid_token.max_budget}"
)
# Token passed all checks # Token passed all checks
# Add token to cache # Add token to cache
user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60) user_api_key_cache.set_cache(key=api_key, value=valid_token, ttl=60)
@ -668,7 +676,9 @@ async def update_database(
if prisma_client is not None: if prisma_client is not None:
# Fetch the existing cost for the given token # Fetch the existing cost for the given token
existing_spend_obj = await prisma_client.get_data(token=token) existing_spend_obj = await prisma_client.get_data(token=token)
verbose_proxy_logger.debug(f"existing spend: {existing_spend_obj}") verbose_proxy_logger.debug(
f"_update_key_db: existing spend: {existing_spend_obj}"
)
if existing_spend_obj is None: if existing_spend_obj is None:
existing_spend = 0 existing_spend = 0
else: else:
@ -679,12 +689,18 @@ async def update_database(
verbose_proxy_logger.debug(f"new cost: {new_spend}") verbose_proxy_logger.debug(f"new cost: {new_spend}")
# Update the cost column for the given token # Update the cost column for the given token
await prisma_client.update_data(token=token, data={"spend": new_spend}) await prisma_client.update_data(token=token, data={"spend": new_spend})
valid_token = user_api_key_cache.get_cache(key=token)
valid_token.spend = new_spend
user_api_key_cache.set_cache(key=token, value=valid_token)
elif custom_db_client is not None: elif custom_db_client is not None:
# Fetch the existing cost for the given token # Fetch the existing cost for the given token
existing_spend_obj = await custom_db_client.get_data( existing_spend_obj = await custom_db_client.get_data(
key=token, table_name="key" key=token, table_name="key"
) )
verbose_proxy_logger.debug(f"existing spend: {existing_spend_obj}") verbose_proxy_logger.debug(
f"_update_key_db existing spend: {existing_spend_obj}"
)
if existing_spend_obj is None: if existing_spend_obj is None:
existing_spend = 0 existing_spend = 0
else: else:
@ -698,6 +714,10 @@ async def update_database(
key=token, value={"spend": new_spend}, table_name="key" key=token, value={"spend": new_spend}, table_name="key"
) )
valid_token = user_api_key_cache.get_cache(key=token)
valid_token.spend = new_spend
user_api_key_cache.set_cache(key=token, value=valid_token)
async def _insert_spend_log_to_db(): async def _insert_spend_log_to_db():
# Helper to generate payload to log # Helper to generate payload to log
verbose_proxy_logger.debug("inserting spend log to db") verbose_proxy_logger.debug("inserting spend log to db")