diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2636d92a7cf..d71d0f8d403 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2284,6 +2284,7 @@ class PrometheusLogger(CustomLogger): _org_alias: Final = _metadata.get("user_api_key_org_alias", None) _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) _committed_at_auth: Final = committed_budget_spend_from_metadata(_carried_metadata) + _row_carries_auth: Final = "user_api_key_auth" in _carried_metadata # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2295,9 +2296,20 @@ class PrometheusLogger(CustomLogger): budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout() async def emit_budget_metrics() -> tuple[object, ...]: + committed_now: Final = await _committed_spend_now( + _committed_at_auth.counter_keys + if _row_carries_auth + else EntitySpendCounterKeys.of( + token=user_api_key, + team_id=user_api_team, + user_id=user_id, + org_id=user_api_key_org_id, + end_user_id=end_user_id, + ) + ) committed: Final = _committed_spend_before_request( - at_auth=_committed_at_auth, - now=await _committed_spend_now(_committed_at_auth.counter_keys), + at_auth=_committed_at_auth if _row_carries_auth else committed_now, + now=committed_now, response_cost=response_cost, ) return await asyncio.gather( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 97018b74b95..7f95121ea23 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -39,6 +39,7 @@ from litellm.proxy.db.db_spend_update_writer import ( from litellm.proxy.hooks.search_free_tier import apply_search_free_tier from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( ObjectMapping, @@ -1022,6 +1023,30 @@ class _IncrementSpendCounters(Protocol): ) -> None: ... +async def _committed_spend_scope_of_row( + kwargs: Mapping[str, object], + response_cost: float, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, +) -> EntitySpendCounterKeys: + metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs))) + if "user_api_key_auth" in metadata: + return committed_budget_spend_from_metadata(metadata).counter_keys + if response_cost <= 0 or not committed_spend_tracked(): + return EntitySpendCounterKeys() + from litellm.proxy.proxy_server import spend_counter_cache + + return await warm_committed_counter_keys( + spend_counter_cache, + EntitySpendCounterKeys.of( + token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id + ), + ) + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -1153,9 +1178,15 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, - committed_spend_scope=committed_budget_spend_from_metadata( - _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=kwargs)) - ).counter_keys, + committed_spend_scope=await _committed_spend_scope_of_row( + kwargs=kwargs, + response_cost=response_cost, + token=user_api_key, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + ), ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py index 671662fa8c0..0f542040d8e 100644 --- a/litellm/proxy/spend_tracking/committed_spend.py +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -108,6 +108,20 @@ async def read_committed_budget_spend( return _committed_budget_spend(counter_keys=counter_keys, committed=committed) +async def warm_committed_counter_keys( + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys +) -> EntitySpendCounterKeys: + """For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists.""" + committed: Final = await read_committed_budget_spend(spend_counter_cache, counter_keys) + return EntitySpendCounterKeys( + key=counter_keys.key if committed.key is not None else None, + team=counter_keys.team if committed.team is not None else None, + user=counter_keys.user if committed.user is not None else None, + org=counter_keys.org if committed.org is not None else None, + end_user=counter_keys.end_user if committed.end_user is not None else None, + ) + + async def read_admission_committed_spend( spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys ) -> CommittedBudgetSpend: diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 74525da4c3e..1a746de588e 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -311,6 +311,44 @@ async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted( } +def _billed_row_metadata() -> dict: + return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"} + + +@pytest.mark.asyncio +async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to( + prometheus_logger, getters, spend_counter_cache +): + """Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(2.0), + "team": pytest.approx(265.0), + "user": pytest.approx(37.0), + "org": 435.0, + } + + +@pytest.mark.asyncio +async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + @pytest.mark.asyncio async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend( prometheus_logger, getters, monkeypatch diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c3f4c4e9d97..c736f92da6d 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1070,6 +1070,54 @@ async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(m assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tracked", "expected_key", "expected_org"), + [(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)], +) +async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist( + tracked, expected_key, expected_org +): + """A background response is billed from its creator's stored attribution, after the request's auth is gone.""" + from litellm.proxy.proxy_server import increment_spend_counters + + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + spend_counter_cache = DualCache() + spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:key:test_api_key", value=1.0) + spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:org:test_org_id", value=2.0) + + with ( + patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked), + patch("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache), + patch( + "litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db", + AsyncMock(return_value=50.0), + ), + ): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + cached = spend_counter_cache.in_memory_cache.get_cache + assert cached(key="committed_spend:key:test_api_key") == expected_key + assert cached(key="committed_spend:org:test_org_id") == expected_org + assert cached(key="committed_spend:team:test_team_id") is None + assert cached(key="committed_spend:user:test_user_id") is None + + @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock()