diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index ae0b39135f1..589d8fe68d1 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -201,11 +201,12 @@ def model_access_group_registry_cache_key() -> str: def model_access_group_spend_counter_key(access_group_name: str) -> str: - """Spend counter key for one model access group; shared so its three owners cannot drift. + """Spend counter key for one model access group; shared so its four owners cannot drift. - The reservation path writes it, auth reads it to enforce ``max_budget``, and the reset job - clears it on rollover. A copy that drifts in any one of them silently resets or reads a - counter nobody else touches, which shows up as a budget that never trips or never resets. + The reservation path writes it up front, the cost callback writes it after the call, auth + reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that + drifts in any one of them silently resets or reads a counter nobody else touches, which shows + up as a budget that never trips or never resets. """ return f"spend:model_access_group:{access_group_name}" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index f593e94b36f..ad857b8c3f5 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,5 +1,6 @@ import asyncio import traceback +from collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -27,6 +28,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, + get_request_model_access_groups, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -258,6 +260,7 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, metadata=metadata, ) + model_access_groups: Final = get_request_model_access_groups(kwargs) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) @@ -296,6 +299,7 @@ class _ProxyDBLogger(CustomLogger): response_cost=response_cost, budget_reservation=budget_reservation, request_tags=tags, + model_access_groups=model_access_groups, ) # update cache (fire-and-forget for backward compat: @@ -572,6 +576,7 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, + model_access_groups: Sequence[str] | None = None, ) -> None: try: await proxy_logging_obj.db_spend_update_writer.update_database( @@ -610,6 +615,7 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, end_user_id=end_user_id, tags=request_tags, + model_access_groups=model_access_groups, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af0aa9743bc..accf03b33ab 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -382,6 +382,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -2648,6 +2650,7 @@ async def increment_spend_counters( budget_reservation: dict | None = None, end_user_id: str | None = None, tags: list[str] | None = None, + model_access_groups: Sequence[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2777,6 +2780,13 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, + _increment_model_access_group_spend_counters( + model_access_groups=model_access_groups, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if model_access_groups + else None, _increment_org_spend_counter( org_id=org_id, response_cost=cost, @@ -2865,6 +2875,33 @@ async def _increment_end_user_and_tag_spend_counters( ) +async def _increment_model_access_group_spend_counters( + model_access_groups: Sequence[object], + response_cost: float, + reserved_counter_keys: set[str], +) -> None: + """Charge the model access groups that authorized this request. + + Without this the counter auth reads is written only by the reservation path, so + ``disable_budget_reservation`` would leave ``_model_access_group_max_budget_check`` enforcing + against the DB row's spend, which lags by up to the cache TTL. + + Typed ``object`` rather than ``str`` because the names reach the cost callback out of request + metadata, which the coercion upstream filters to a list but not to strings. A non-string that + slipped through would build a counter key nothing else ever reads. + """ + unique_groups: Final = tuple( + dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) + ) + for group in unique_groups: + await _init_and_increment_unreserved_spend_counter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + + async def _increment_org_spend_counter( org_id: str | None, response_cost: float, diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index ca517474a5c..e923a28f371 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,14 +1,14 @@ -import pytest - - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _ProxyDBLogger, _get_budget_reservation_from_metadata, + _ProxyDBLogger, _should_track_cost_callback, _update_database_and_spend_counters, ) @@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda response_cost=0.2, budget_reservation=budget_reservation, request_tags=["tag-a"], + model_access_groups=("premium",), ) proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() @@ -598,6 +599,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda budget_reservation=budget_reservation, end_user_id="test_end_user_id", tags=["tag-a"], + model_access_groups=("premium",), ) @@ -1875,3 +1877,82 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( 1 if expect_spend_log else 0 ) + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped(): + """Auth stamps the matched groups onto request metadata; the callback has to carry them through. + + Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with + reservations disabled the budget check reads a counter no one maintains. + """ + logger = _ProxyDBLogger() + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_user_id": "user-1", + MODEL_ACCESS_GROUP_METADATA_KEY: ["premium", "starter"], + }, + }, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, + "stream": False, + } + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment_spend_counters, + patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_increment_spend_counters.assert_awaited_once() + assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == ( + "premium", + "starter", + ) + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped(): + """A request no budgeted group authorized must not debit anything.""" + logger = _ProxyDBLogger() + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "litellm_params": { + "metadata": {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"}, + }, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None}, + "stream": False, + } + + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment_spend_counters, + patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_increment_spend_counters.assert_awaited_once() + assert mock_increment_spend_counters.await_args.kwargs["model_access_groups"] == () diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1b28566366c..65c8c248a27 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -33,6 +33,7 @@ from litellm.proxy.common_utils.reset_budget_job import _model_access_group_coun from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, model_access_group_cache_key, + model_access_group_spend_counter_key, ) from litellm.proxy.spend_tracking.budget_reservation import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, @@ -47,6 +48,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @pytest.fixture() @@ -3067,3 +3069,117 @@ async def test_model_access_group_counter_blocks_a_request_over_the_group_budget assert exc_info.value.entity_id == "premium" assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + + +async def _cache_model_access_group_budget(key_cache, group, spend, max_budget=None): + await key_cache.async_set_cache( + key=model_access_group_cache_key(group), + value=ModelAccessGroupBudget(access_group_name=group, spend=spend, max_budget=max_budget), + model_type=ModelAccessGroupBudget, + ) + + +async def _reserve_for_model_access_groups(key_cache, groups, estimate): + """Reserve against the given groups, whose rows are already cached, so nothing hits the DB.""" + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=estimate, + ): + return await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth( + api_key="hashed", token="tok-mag-counter", matched_model_access_groups=list(groups) + ), + team_object=None, + user_object=None, + prisma_client=_ModelAccessGroupBudgetPrisma(), + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_counter_accumulates_across_calls_without_a_reservation(spend_counter_state): + """With reservations disabled nothing writes the counter up front, so the cost callback must. + + Otherwise the read-time budget check enforces against the DB row's spend, which the cache + holds for the full TTL, and a caller runs past the ceiling for that whole window. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + from litellm.proxy.proxy_server import increment_spend_counters + + counter_key = model_access_group_spend_counter_key("premium") + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.25, model_access_groups=["premium"] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.25) + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.75, model_access_groups=["premium", "premium", ""] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.0) + assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None + + +@pytest.mark.asyncio +async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state): + """The reservation already wrote this counter, so the post-call pass has to skip it.""" + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium"], estimate=0.6) + counter_key = model_access_group_spend_counter_key("premium") + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.6) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium"], + ) + + # 1.0 recorded + the reservation reconciled down to the 0.2 actually spent. A second + # increment would land at 1.4. + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.2) + + +@pytest.mark.asyncio +async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one(spend_counter_state): + """A budgetless group reserves nothing, so only the post-call pass can charge it. + + Both groups authorized the request and both get debited, each exactly once, whether or not + the reservation path happened to hold a counter for them. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + await _cache_model_access_group_budget(key_cache, "starter", spend=4.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium", "starter"], estimate=0.6) + assert [entry["entity_id"] for entry in reservation["entries"]] == ["premium"] + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium", "starter", "starter", "premium"], + ) + + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("premium") + ) == pytest.approx(1.2) + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("starter") + ) == pytest.approx(4.2)