scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fallback to plain org cache for spend counters

Browse source
This commit is contained in:
user 2026-04-30 20:04:16 -07:00
parent c28e093f41
commit dcfde1b899
2 changed files with 54 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/proxy/proxy_server.py
View file

@ -2048,7 +2048,7 @@ async def _increment_org_spend_counter(
await _init_and_increment_unreserved_spend_counter(
counter_key=f"spend:org:{org_id}",
source_cache_key=f"org_id:{org_id}:with_budget",
source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
@ -2056,7 +2056,7 @@ async def _increment_org_spend_counter(
async def _init_and_increment_unreserved_spend_counter(
counter_key: str,
source_cache_key: str,
source_cache_key: Union[str, List[str]],
increment: float,
reserved_counter_keys: Set[str],
) -> None:
@ -2072,7 +2072,7 @@ async def _init_and_increment_unreserved_spend_counter(
async def _init_and_increment_spend_counter(
counter_key: str,
source_cache_key: str,
source_cache_key: Union[str, List[str]],
increment: float,
):
"""
@ -2117,7 +2117,7 @@ async def _init_and_increment_window_spend_counter(
async def _ensure_spend_counter_initialized(
counter_key: str,
source_cache_key: str,
source_cache_key: Union[str, List[str]],
):
is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key)
if is_warm is False:
@ -2130,19 +2130,31 @@ async def _ensure_spend_counter_initialized(
)
if db_spend is None:
# DB unavailable - fall back to in-process cache (may be stale).
source = await user_api_key_cache.async_get_cache(key=source_cache_key)
base_spend: float = 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
base_spend = await _get_source_cache_base_spend(
source_cache_key=source_cache_key
)
if base_spend > 0:
await _increment_spend_counter_cache(
counter_key=counter_key, increment=base_spend
)
async def _get_source_cache_base_spend(
source_cache_key: Union[str, List[str]],
) -> float:
source_cache_keys = (
[source_cache_key] if isinstance(source_cache_key, str) else source_cache_key
)
for cache_key in source_cache_keys:
source = await user_api_key_cache.async_get_cache(key=cache_key)
if source is None:
continue
if isinstance(source, dict):
return float(source.get("spend", 0.0) or 0.0)
return float(getattr(source, "spend", 0.0) or 0.0)
return 0.0
async def _ensure_window_spend_counter_initialized(
counter_key: str,
entity_type: str,

31
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -509,6 +509,37 @@ async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_stat
) == pytest.approx(2.25)
@pytest.mark.asyncio
async def test_should_seed_org_counter_from_plain_org_cache(spend_counter_state):
counter_cache, key_cache = spend_counter_state
await key_cache.async_set_cache(
key="org_id:org-counter-plain",
value=LiteLLM_OrganizationTable(
organization_id="org-counter-plain",
organization_alias="shared-org",
budget_id="org-budget-id",
spend=2.0,
models=[],
created_by="test",
updated_by="test",
).model_dump(),
)
from litellm.proxy.proxy_server import increment_spend_counters
await increment_spend_counters(
token=None,
team_id=None,
user_id=None,
org_id="org-counter-plain",
response_cost=0.25,
)
assert counter_cache.in_memory_cache.get_cache(
key="spend:org:org-counter-plain"
) == pytest.approx(2.25)
@pytest.mark.asyncio
async def test_should_cap_known_estimate_to_remaining_budget(
spend_counter_state,