scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(proxy/auth): handle tz-aware temp_budget_expiry (#33840)

Browse source 
Co-authored-by: shivam <shivam@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
(cherry picked from commit 10d2a27d87)
This commit is contained in:
devin-ai-integration[bot] 2026-07-20 17:43:57 -07:00 • committed by Yuneng Jiang
parent 13fef68b12
commit 2374dda83e
No known key found for this signature in database
2 changed files with 32 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/auth/user_api_key_auth.py
View file

@ -2685,7 +2685,9 @@ def _get_temp_budget_increase(valid_token: UserAPIKeyAuth):
valid_token_metadata = valid_token.metadata
if "temp_budget_increase" in valid_token_metadata and "temp_budget_expiry" in valid_token_metadata:
expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"])
if expiry > datetime.now():
if expiry.tzinfo is None:
expiry = expiry.replace(tzinfo=timezone.utc)
if expiry > datetime.now(timezone.utc):
return valid_token_metadata["temp_budget_increase"]
return None

29
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -1732,6 +1732,35 @@ def test_get_temp_budget_increase():
assert _get_temp_budget_increase(valid_token) == 100
def test_get_temp_budget_increase_tz_aware_expiry():
from datetime import datetime, timedelta, timezone
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import _get_temp_budget_increase
future_expiry = (datetime.now(timezone.utc) + timedelta(days=1)).isoformat()
valid_token = UserAPIKeyAuth(
max_budget=100,
spend=0,
metadata={
"temp_budget_increase": 100,
"temp_budget_expiry": future_expiry,
},
)
assert _get_temp_budget_increase(valid_token) == 100
past_expiry = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat()
expired_token = UserAPIKeyAuth(
max_budget=100,
spend=0,
metadata={
"temp_budget_increase": 100,
"temp_budget_expiry": past_expiry,
},
)
assert _get_temp_budget_increase(expired_token) is None
def test_update_key_budget_with_temp_budget_increase():
from datetime import datetime, timedelta