mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(proxy): give the team member spend reset the spend counter Redis TTL
This commit is contained in:
parent
1985f78908
commit
21aa23caae
2 changed files with 12 additions and 3 deletions
|
|
@ -3198,7 +3198,11 @@ async def invalidate_team_member_spend_state(
|
|||
await asyncio.wait((inflight,))
|
||||
|
||||
if new_spend is not None:
|
||||
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
|
||||
from litellm.proxy.proxy_server import (
|
||||
SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
|
||||
spend_counter_cache,
|
||||
spend_counter_redis_ttl,
|
||||
)
|
||||
|
||||
spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
|
||||
spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
|
||||
|
|
@ -3206,7 +3210,9 @@ async def invalidate_team_member_spend_state(
|
|||
spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
|
||||
await spend_counter_cache.redis_cache.async_set_cache(
|
||||
key=spend_counter_key, value=new_spend, ttl=spend_counter_redis_ttl()
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
|
||||
|
|
|
|||
|
|
@ -10011,6 +10011,7 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
|
|||
caller passing new_spend must match that precedent, not merely delete the counter."""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.proxy_server import spend_counter_redis_ttl
|
||||
|
||||
real_cache = UserApiKeyCache()
|
||||
real_spend_counter_cache = DualCache()
|
||||
|
|
@ -10029,7 +10030,9 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele
|
|||
)
|
||||
|
||||
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
|
||||
fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
|
||||
fake_redis_cache.async_set_cache.assert_awaited_once_with(
|
||||
key="spend:team_member:user-1:team-1", value=2.5, ttl=spend_counter_redis_ttl()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue