scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): enforce budget limits across multi-pod deployments via Redis-backed spend counters

Browse source 
Budget checks on API keys, teams, and team members were not enforced in
multi-pod deployments because user_api_key_cache is intentionally
in-memory-only. Each pod tracked spend independently, so with N pods
the effective budget was N × max_budget.

Introduces a separate spend_counter_cache (DualCache wired to
redis_usage_cache) with atomic increment/read helpers:
- increment_spend_counters(): awaited in cost callback (not create_task)
  to update both in-memory and Redis before the next auth check
- get_current_spend(): reads Redis first (cross-pod authoritative),
  falls back to in-memory, then to cached object .spend from DB

Budget check functions (_virtual_key_max_budget_check,
_team_max_budget_check, _check_team_member_budget) now read spend via
get_current_spend() instead of cached object .spend fields.

When Redis is not configured, falls back to in-memory-only counters
(same as current single-instance behavior).

Fixes #23714
This commit is contained in:
michelligabriele 2026-03-27 14:42:05 +01:00
parent ff63df25a2
commit d533b432fd
No known key found for this signature in database
7 changed files with 618 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

77
litellm/proxy/auth/auth_checks.py
View file

@ -2876,7 +2876,15 @@ async def _virtual_key_max_budget_check(
Triggers a budget alert if the token is over it's max budget.
"""
if valid_token.spend is not None and valid_token.max_budget is not None:
if valid_token.max_budget is not None:
from litellm.proxy.proxy_server import get_current_spend
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
spend = await get_current_spend(
counter_key=f"spend:key:{valid_token.token}",
fallback_spend=valid_token.spend or 0.0,
)
####################################
# collect information for alerting #
####################################
@ -2888,7 +2896,7 @@ async def _virtual_key_max_budget_check(
call_info = CallInfo(
token=valid_token.token,
spend=valid_token.spend,
spend=spend,
max_budget=valid_token.max_budget,
soft_budget=valid_token.soft_budget,
user_id=valid_token.user_id,
@ -2909,9 +2917,9 @@ async def _virtual_key_max_budget_check(
# collect information for alerting #
####################################
if valid_token.spend >= valid_token.max_budget:
if spend >= valid_token.max_budget:
raise litellm.BudgetExceededError(
current_cost=valid_token.spend,
current_cost=spend,
max_budget=valid_token.max_budget,
)
@ -3042,6 +3050,14 @@ async def _check_team_member_budget(
team_member_budget = team_membership.litellm_budget_table.max_budget
team_member_spend = team_membership.spend or 0.0
# Read from cross-pod counter (Redis-first) if available
from litellm.proxy.proxy_server import get_current_spend
team_member_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
fallback_spend=team_member_spend,
)
if team_member_spend >= team_member_budget:
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
@ -3065,33 +3081,40 @@ async def _team_max_budget_check(
if (
team_object is not None
and team_object.max_budget is not None
and team_object.spend is not None
and team_object.spend > team_object.max_budget
):
if valid_token:
call_info = CallInfo(
token=valid_token.token,
spend=team_object.spend,
max_budget=team_object.max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=valid_token.org_id,
event_group=Litellm_EntityType.TEAM,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="team_budget",
user_info=call_info,
)
)
from litellm.proxy.proxy_server import get_current_spend
raise litellm.BudgetExceededError(
current_cost=team_object.spend,
max_budget=team_object.max_budget,
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}",
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
spend = await get_current_spend(
counter_key=f"spend:team:{team_object.team_id}",
fallback_spend=team_object.spend or 0.0,
)
if spend > team_object.max_budget:
if valid_token:
call_info = CallInfo(
token=valid_token.token,
spend=spend,
max_budget=team_object.max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=valid_token.org_id,
event_group=Litellm_EntityType.TEAM,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="team_budget",
user_info=call_info,
)
)
raise litellm.BudgetExceededError(
current_cost=spend,
max_budget=team_object.max_budget,
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}",
)
async def _team_soft_budget_check(
team_object: Optional[LiteLLM_TeamTable],

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

@ -1303,9 +1303,21 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
team_member_info.litellm_budget_table.max_budget
)
if team_member_budget is not None and team_member_budget > 0:
if valid_token.team_member_spend > team_member_budget:
# Read from cross-pod counter (Redis-first) if available
from litellm.proxy.proxy_server import get_current_spend
team_member_spend = valid_token.team_member_spend
if (
valid_token.user_id is not None
and valid_token.team_id is not None
):
team_member_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}",
fallback_spend=team_member_spend,
)
if team_member_spend > team_member_budget:
raise litellm.BudgetExceededError(
current_cost=valid_token.team_member_spend,
current_cost=team_member_spend,
max_budget=team_member_budget,
)

84
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -54,16 +54,49 @@ class ResetBudgetJob:
"""
Resets the budget for all LiteLLM Team Members if their budget has expired
"""
budget_ids = [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
# Reset spend counters for affected team members.
# Reset Redis directly so a transient failure doesn't leave stale
# counters that get_current_spend would read as authoritative.
try:
from litellm.proxy.proxy_server import spend_counter_cache
memberships = (
await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
)
for m in memberships:
counter_key = f"spend:team_member:{m.user_id}:{m.team_id}"
# Always reset in-memory
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0
)
# Explicitly reset Redis with warning on failure
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset team member spend counter in Redis %s: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to reset team member spend counters: %s", e
)
return await self.prisma_client.db.litellm_teammembership.update_many(
where={
"budget_id": {
"in": [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
}
},
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
@ -531,6 +564,39 @@ class ResetBudgetJob:
"""
try:
item.spend = 0.0
# Reset the cross-pod spend counter.
# Reset Redis directly (not via DualCache) so a Redis failure
# doesn't silently leave a stale counter that get_current_spend
# would read as authoritative, permanently blocking the user.
from litellm.proxy.proxy_server import spend_counter_cache
counter_key = None
if item_type == "key" and hasattr(item, "token") and item.token is not None:
counter_key = f"spend:key:{item.token}"
elif item_type == "team" and hasattr(item, "team_id") and item.team_id is not None:
counter_key = f"spend:team:{item.team_id}"
if counter_key is not None:
# Always reset in-memory (local fallback)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0
)
# Explicitly reset Redis with warning on failure
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset spend counter in Redis for %s key=%s: %s. "
"Budget may be over-enforced until counter expires.",
item_type,
counter_key,
redis_err,
)
if hasattr(item, "budget_duration") and item.budget_duration is not None:
# Get standardized reset time based on budget duration
from litellm.proxy.common_utils.timezone_utils import (

18
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -135,7 +135,11 @@ class _ProxyDBLogger(CustomLogger):
start_time=None,
end_time=None, # start/end time for completion
):
from litellm.proxy.proxy_server import proxy_logging_obj, update_cache
from litellm.proxy.proxy_server import (
increment_spend_counters,
proxy_logging_obj,
update_cache,
)
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
try:
@ -194,7 +198,17 @@ class _ProxyDBLogger(CustomLogger):
org_id=org_id,
)
# update cache
# Atomically update spend counters (in-memory + Redis)
# for cross-pod budget enforcement.
await increment_spend_counters(
token=user_api_key,
team_id=team_id,
user_id=user_id,
response_cost=response_cost,
)
# update cache (fire-and-forget for backward compat:
# cached object fields, soft budget alerts, etc.)
asyncio.create_task(
update_cache(
token=user_api_key,

132
litellm/proxy/proxy_server.py
View file

@ -1530,6 +1530,9 @@ shared_aiohttp_session: Optional[
user_api_key_cache = DualCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
spend_counter_cache = DualCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
dual_cache=user_api_key_cache
)
@ -1694,6 +1697,134 @@ def cost_tracking():
)
async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
"""
Read current spend from the cross-pod spend counter.
Reads Redis FIRST (authoritative cross-pod value), not DualCache's
async_get_cache which returns in-memory first. This is critical:
DualCache.async_get_cache returns stale per-pod values because each
pod's in-memory cache is only updated by that pod's own increments.
Fallback chain:
1. Redis counter (cross-pod, authoritative)
2. In-memory counter (single-instance or Redis failure)
3. Cached object's .spend from DB (cold start, no counter yet)
"""
# 1. Try Redis first (cross-pod authoritative)
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(
key=counter_key
)
if val is not None:
return float(val)
except Exception as e:
verbose_proxy_logger.debug(
"get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
counter_key,
e,
)
# 2. Fall back to in-memory counter (single-instance or Redis failure)
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
# 3. Final fallback: cached object's spend from DB
return fallback_spend
async def increment_spend_counters(
token: Optional[str],
team_id: Optional[str],
user_id: Optional[str],
response_cost: Optional[float],
):
"""
Atomically increment spend counters for budget enforcement.
Uses spend_counter_cache (DualCache with Redis backend when available)
so counters are shared across all pods. Budget check functions read
from these counters via get_current_spend() (Redis-first).
Awaited (not create_task) in the cost callback, so the counter is
updated before the next request's auth check runs.
"""
if response_cost is None or response_cost == 0:
return
if token is not None:
# token arrives pre-hashed from metadata["user_api_key"] (auth flow
# hashes raw "sk-..." keys before they reach the callback). The
# startswith("sk-") check is a safety net matching update_cache —
# if a raw key somehow arrives, hash it; otherwise use as-is to
# avoid double-hashing (budget checks read valid_token.token which
# is single-hashed).
hashed_token = (
hash_token(token=token)
if isinstance(token, str) and token.startswith("sk-")
else token
)
await _init_and_increment_spend_counter(
counter_key=f"spend:key:{hashed_token}",
source_cache_key=hashed_token,
increment=response_cost,
)
if team_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:team:{team_id}",
source_cache_key=f"team_id:{team_id}",
increment=response_cost,
)
if user_id is not None and team_id is not None:
await _init_and_increment_spend_counter(
counter_key=f"spend:team_member:{user_id}:{team_id}",
source_cache_key=f"team_membership:{user_id}:{team_id}",
increment=response_cost,
)
async def _init_and_increment_spend_counter(
counter_key: str,
source_cache_key: str,
increment: float,
):
"""
Initialize counter from cached object's DB-loaded spend if not yet set,
then atomically increment in both in-memory and Redis.
On first access per pod:
1. Check spend_counter_cache (in-memory -> Redis via DualCache for init check)
2. If not found anywhere, read base spend from user_api_key_cache (DB-loaded object)
3. Seed counter via async_increment_cache (not async_set_cache) to avoid a
check-then-set race: if two pods cold-start simultaneously, both may see
the counter as absent and seed it. Using increment instead of set means
the worst case is over-counting (conservative — blocks slightly early)
rather than under-counting (would allow overspend).
4. Increment atomically (both in-memory + Redis)
"""
current = await spend_counter_cache.async_get_cache(key=counter_key)
if current is None:
source = await user_api_key_cache.async_get_cache(key=source_cache_key)
base_spend = 0.0
if source is not None:
if isinstance(source, dict):
base_spend = source.get("spend", 0.0) or 0.0
else:
base_spend = getattr(source, "spend", 0.0) or 0.0
if base_spend > 0:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=base_spend
)
await spend_counter_cache.async_increment_cache(
key=counter_key, value=increment
)
async def update_cache( # noqa: PLR0915
token: Optional[str],
user_id: Optional[str],
@ -2528,6 +2659,7 @@ class ProxyConfig:
):
## INIT PROXY REDIS USAGE CLIENT ##
redis_usage_cache = litellm.cache.cache
spend_counter_cache.redis_cache = redis_usage_cache
# Note: PKCE verifier storage uses redis_usage_cache directly (not
# user_api_key_cache) to avoid routing all API-key lookups through Redis.

151
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -29,10 +29,13 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
_can_object_call_vector_stores,
_check_team_member_budget,
_get_fuzzy_user_object,
_get_team_db_check,
_log_budget_lookup_failure,
_team_max_budget_check,
_virtual_key_max_budget_alert_check,
_virtual_key_max_budget_check,
_virtual_key_soft_budget_check,
get_key_object,
get_user_object,
@ -1629,3 +1632,151 @@ async def test_custom_auth_common_checks_opt_in():
parent_otel_span=None,
)
mock_common.assert_called_once()
# =====================================================================
# Spend counter budget check tests (v2 — Redis-backed spend counters)
# =====================================================================
@pytest.mark.asyncio
async def test_virtual_key_budget_check_reads_from_spend_counter():
"""Budget check should use get_current_spend when counter exists,
even if cached object shows lower spend."""
from litellm.proxy.utils import ProxyLogging
valid_token = UserAPIKeyAuth(
token="test-hashed-token",
spend=0.0, # stale — counter has 1.5
max_budget=1.0,
user_id="test-user",
)
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
proxy_logging_obj.budget_alerts = AsyncMock()
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:key:test-hashed-token":
return 1.5
return fallback_spend
with patch(
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
assert exc_info.value.current_cost == 1.5
assert exc_info.value.max_budget == 1.0
@pytest.mark.asyncio
async def test_virtual_key_budget_check_fallback_no_counter():
"""When counter doesn't exist, budget check should fall back
to cached object's spend via fallback_spend."""
from litellm.proxy.utils import ProxyLogging
valid_token = UserAPIKeyAuth(
token="test-hashed-token",
spend=15.0,
max_budget=10.0,
user_id="test-user",
)
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
proxy_logging_obj.budget_alerts = AsyncMock()
# get_current_spend returns fallback_spend when no counter exists
async def mock_get_current_spend(counter_key, fallback_spend):
return fallback_spend
with patch(
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
assert exc_info.value.current_cost == 15.0
@pytest.mark.asyncio
async def test_team_budget_check_reads_from_spend_counter():
"""Team budget check should use get_current_spend when counter exists."""
from litellm.proxy.utils import ProxyLogging
team_object = LiteLLM_TeamTable(
team_id="test-team",
spend=0.0, # stale
max_budget=1.0,
)
valid_token = UserAPIKeyAuth(token="test-token", team_id="test-team")
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
proxy_logging_obj.budget_alerts = AsyncMock()
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:team:test-team":
return 1.5
return fallback_spend
with patch(
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _team_max_budget_check(
team_object=team_object,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
assert exc_info.value.current_cost == 1.5
@pytest.mark.asyncio
async def test_team_member_budget_check_reads_from_spend_counter():
"""Team member budget check should use get_current_spend when counter exists."""
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership
from litellm.proxy.utils import ProxyLogging
team_object = LiteLLM_TeamTable(team_id="test-team")
user_object = LiteLLM_UserTable(user_id="test-user")
valid_token = UserAPIKeyAuth(
token="test-token",
user_id="test-user",
team_id="test-team",
)
team_membership = LiteLLM_TeamMembership(
user_id="test-user",
team_id="test-team",
spend=0.0, # stale
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
)
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
async def mock_get_current_spend(counter_key, fallback_spend):
if counter_key == "spend:team_member:test-user:test-team":
return 1.5
return fallback_spend
with patch(
"litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend
), patch(
"litellm.proxy.auth.auth_checks.get_team_membership",
new_callable=AsyncMock,
return_value=team_membership,
):
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _check_team_member_budget(
team_object=team_object,
user_object=user_object,
valid_token=valid_token,
prisma_client=MagicMock(),
user_api_key_cache=MagicMock(),
proxy_logging_obj=proxy_logging_obj,
)
assert exc_info.value.current_cost == 1.5

180
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -4352,3 +4352,183 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch):
# add_deployment should NOT have been called since store_model_in_db is False
mock_proxy_config.add_deployment.assert_not_called()
# =====================================================================
# Spend counter tests (v2 — Redis-backed spend counters)
# =====================================================================
@pytest.mark.asyncio
async def test_get_current_spend_reads_redis_first():
"""get_current_spend should prefer Redis over in-memory."""
from litellm.caching.dual_cache import DualCache
counter_cache = DualCache()
# In-memory has stale value
counter_cache.in_memory_cache.set_cache(key="spend:key:test", value=0.30)
# Mock Redis with cross-pod authoritative value
mock_redis = AsyncMock()
mock_redis.async_get_cache = AsyncMock(return_value=0.90)
counter_cache.redis_cache = mock_redis
import litellm.proxy.proxy_server as ps
original = ps.spend_counter_cache
ps.spend_counter_cache = counter_cache
try:
from litellm.proxy.proxy_server import get_current_spend
result = await get_current_spend(
counter_key="spend:key:test",
fallback_spend=0.0,
)
# Should return Redis value (0.90), not in-memory (0.30)
assert result == 0.90
mock_redis.async_get_cache.assert_called_once_with(key="spend:key:test")
finally:
ps.spend_counter_cache = original
@pytest.mark.asyncio
async def test_get_current_spend_fallback_to_in_memory():
"""When Redis is not configured, get_current_spend uses in-memory."""
from litellm.caching.dual_cache import DualCache
counter_cache = DualCache() # no redis_cache
counter_cache.in_memory_cache.set_cache(key="spend:key:test", value=0.50)
import litellm.proxy.proxy_server as ps
original = ps.spend_counter_cache
ps.spend_counter_cache = counter_cache
try:
from litellm.proxy.proxy_server import get_current_spend
result = await get_current_spend(
counter_key="spend:key:test",
fallback_spend=0.0,
)
assert result == 0.50
finally:
ps.spend_counter_cache = original
@pytest.mark.asyncio
async def test_increment_spend_counters_initializes_and_increments():
"""Counter should initialize from cached object spend, then increment.
Uses a pre-hashed token to match production: metadata["user_api_key"]
is always hashed by the auth flow before reaching the cost callback.
"""
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_VerificationTokenView, hash_token
key_cache = DualCache()
counter_cache = DualCache()
# In production, the auth flow hashes the raw key before it reaches
# the cost callback. Simulate that by passing the hashed token.
hashed_token = hash_token("sk-test-token-for-counter")
# Simulate a cached key object with existing spend from DB
cached_key = LiteLLM_VerificationTokenView(
token=hashed_token,
spend=5.0,
max_budget=10.0,
)
key_cache.in_memory_cache.set_cache(key=hashed_token, value=cached_key)
import litellm.proxy.proxy_server as ps
original_key_cache = ps.user_api_key_cache
original_counter_cache = ps.spend_counter_cache
ps.user_api_key_cache = key_cache
ps.spend_counter_cache = counter_cache
try:
from litellm.proxy.proxy_server import increment_spend_counters
# Pass pre-hashed token (as the cost callback would in production)
await increment_spend_counters(
token=hashed_token,
team_id=None,
user_id=None,
response_cost=0.50,
)
# Counter should be: base(5.0) + increment(0.50) = 5.50
counter = counter_cache.in_memory_cache.get_cache(
key=f"spend:key:{hashed_token}"
)
assert counter == 5.50
# Second increment — counter already exists, just increment
await increment_spend_counters(
token=hashed_token,
team_id=None,
user_id=None,
response_cost=0.25,
)
counter = counter_cache.in_memory_cache.get_cache(
key=f"spend:key:{hashed_token}"
)
assert counter == 5.75
finally:
ps.user_api_key_cache = original_key_cache
ps.spend_counter_cache = original_counter_cache
@pytest.mark.asyncio
async def test_increment_spend_counters_team_and_member():
"""Counter should track team and team member spend separately."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import LiteLLM_TeamTable
key_cache = DualCache()
counter_cache = DualCache()
# Cached team object
team_obj = LiteLLM_TeamTable(team_id="team-1", spend=2.0)
key_cache.in_memory_cache.set_cache(key="team_id:team-1", value=team_obj)
# Cached team membership
key_cache.in_memory_cache.set_cache(
key="team_membership:user-1:team-1",
value={"user_id": "user-1", "team_id": "team-1", "spend": 1.0},
)
import litellm.proxy.proxy_server as ps
original_key_cache = ps.user_api_key_cache
original_counter_cache = ps.spend_counter_cache
ps.user_api_key_cache = key_cache
ps.spend_counter_cache = counter_cache
try:
from litellm.proxy.proxy_server import increment_spend_counters
await increment_spend_counters(
token=None,
team_id="team-1",
user_id="user-1",
response_cost=0.30,
)
team_counter = counter_cache.in_memory_cache.get_cache(
key="spend:team:team-1"
)
assert team_counter == 2.30
member_counter = counter_cache.in_memory_cache.get_cache(
key="spend:team_member:user-1:team-1"
)
assert member_counter == 1.30
finally:
ps.user_api_key_cache = original_key_cache
ps.spend_counter_cache = original_counter_cache