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

fix(parallel_request_limiter.py): fix team rate limit enforcement

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-26 18:06:13 -08:00
parent f84ac35000
commit f86ab19067
3 changed files with 105 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/proxy_server.py
View file

@ -350,8 +350,7 @@ async def user_api_key_auth(
original_api_key = api_key # (Patch: For DynamoDB Backwards Compatibility)
if api_key.startswith("sk-"):
api_key = hash_token(token=api_key)
# valid_token = user_api_key_cache.get_cache(key=api_key)
valid_token = None
valid_token = user_api_key_cache.get_cache(key=api_key)
if valid_token is None:
## check db
verbose_proxy_logger.debug(f"api key: {api_key}")
@ -384,7 +383,6 @@ async def user_api_key_auth(
# 6. If token spend per model is under budget per model
# 7. If token spend is under team budget
# 8. If team spend is under team budget
request_data = await _read_request_body(
request=request
) # request data, used across all checks. Making this easily available
@ -627,7 +625,7 @@ async def user_api_key_auth(
)
)
if valid_token.spend > valid_token.team_max_budget:
if valid_token.spend >= valid_token.team_max_budget:
raise Exception(
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Team: {valid_token.team_max_budget}"
)
@ -646,7 +644,7 @@ async def user_api_key_auth(
)
)
if valid_token.team_spend > valid_token.team_max_budget:
if valid_token.team_spend >= valid_token.team_max_budget:
raise Exception(
f"ExceededTokenBudget: Current Team Spend: {valid_token.team_spend}; Max Budget for Team: {valid_token.team_max_budget}"
)