diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 352a166fae0..c711f8b5a62 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -180,7 +180,7 @@ async def _apply_over_budget_reservation_policy( entry: dict[str, float | str], applied_entries: list[dict[str, float | str]], reservation_cost: float, - current_spend: float, + spend_before_reservation: float, fail_closed_budget_enforcement: bool = False, ) -> float: """ @@ -196,13 +196,13 @@ async def _apply_over_budget_reservation_policy( applied_entries.remove(entry) return reservation_cost - remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) + remaining_before_reservation: Final = counter.max_budget - spend_before_reservation if remaining_before_reservation <= 1e-12: - _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) - if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: + _raise_counter_budget_exceeded(counter=counter, current_cost=spend_before_reservation + reservation_cost) + if fail_closed_budget_enforcement and reservation_cost - remaining_before_reservation > 1e-12: _raise_counter_budget_exceeded( counter=counter, - current_cost=current_spend - reservation_cost, + current_cost=spend_before_reservation, estimated_cost=reservation_cost, ) await _resize_applied_reservation( @@ -992,7 +992,7 @@ async def _reserve_reservable_counters( current_spend_by_counter_key=current_spend_by_counter_key, reservation_cost=reservation_cost, ): - charged_cost = reservation_cost + charged_cost = min(reservation_cost, max(counter.max_budget for counter in group)) entries = tuple(_counter_to_reservation_entry(counter=counter, reserved_cost=charged_cost) for counter in group) applied_entries.extend(entries) reserved_values = await _reserve_counters(counters=group, entries=entries, reservation_cost=charged_cost) @@ -1005,18 +1005,19 @@ async def _reserve_reservable_counters( for counter, entry, reserved_value in zip(group, entries, reserved_values): if entry not in applied_entries: continue - if reserved_value is not None: - current_spend = reserved_value - (charged_cost - reservation_cost) - else: - current_spend = current_spend_by_counter_key[counter.counter_key] + reservation_cost - if current_spend > counter.max_budget: + spend_before_reservation = ( + current_spend_by_counter_key[counter.counter_key] + if reserved_value is None + else reserved_value - charged_cost + ) + if spend_before_reservation + reservation_cost > counter.max_budget: reservation_cost = await _apply_over_budget_reservation_policy( counter=counter, valid_token=valid_token, entry=entry, applied_entries=applied_entries, reservation_cost=reservation_cost, - current_spend=current_spend, + spend_before_reservation=spend_before_reservation, fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) return reservation_cost diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a719dbca091..744d20ab1c7 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -357,6 +357,31 @@ def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None: assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold +HOLDS_FAR_OVER_A_TINY_BUDGET: Final = ( + (GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}), + ("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("route", "request_body"), HOLDS_FAR_OVER_A_TINY_BUDGET) +async def test_a_hold_far_over_the_budget_keeps_the_spend_already_counted( + spend_counter_cache: DualCache, route: str, request_body: dict[str, object] +) -> None: + counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" + counted_spend: Final = 0.004 + hold_estimate: Final = estimate_request_max_cost(request_body=request_body, route=route, llm_router=None) + assert hold_estimate is not None and 1e20 < hold_estimate < math.inf + await spend_counter_cache.async_set_cache(key=counter_key, value=counted_spend) + + reservation: Final = await _reserve_for_tiny_budget_key(route, request_body) + + assert reservation is not None + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.01) + await release_unbound_budget_reservation(reservation) + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(counted_spend) + + def test_a_negative_max_tokens_holds_the_default_output() -> None: uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY) assert uncapped_hold is not None diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 535c3da8c3a..25e0ee6a0ea 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -207,7 +207,7 @@ async def test_over_budget_window_counter_tags_clean_entity_id(): entry={"counter_key": counter.counter_key}, applied_entries=[], reservation_cost=0.5, - current_spend=2.0, + spend_before_reservation=1.5, ) assert exc_info.value.entity_type == "key" assert exc_info.value.entity_id == "test-token"