scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix(proxy): reserve per-image cost for image-generation requests

Browse source 
Image-generation routes (dall-e-3, flux, etc.) have no per-token output
cost so they fell through to the no-reservation read-time-only path.
Concurrent image requests against a depleted budget could all pass
common_checks (counter exactly at max_budget passes the strict-`>`
gate) and reach the provider before reconciliation caught up.

Add per-image reservation in _estimate_request_max_cost_for_model:
when the model has a per-image cost field, reserve `n × cost_per_image`
upfront. The atomic counter increment serializes concurrent admissions,
so the second request sees the post-first-reservation counter and
raises BudgetExceededError instead of silently leaking through.

Both `output_cost_per_image` and `input_cost_per_image` are honored —
naming is inconsistent across providers (OpenAI dall-e-3 uses
input_cost_per_image, aiml/dall-e-3 uses output_cost_per_image for
the same per-generated-image price).

Per-pixel pricing (DALL-E 2 size variants) and TTS/STT routes still
fall through to read-time enforcement; those are follow-ups.
This commit is contained in:
Yuneng Jiang 2026-05-08 21:08:55 -07:00
parent 4901ecc6b8
commit 0d551ac4f0
No known key found for this signature in database
2 changed files with 189 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -827,6 +827,13 @@ def _estimate_request_max_cost_for_model(
if model_info is None:
return None
image_cost = _estimate_image_generation_cost(
request_body=request_body,
model_info=model_info,
)
if image_cost is not None:
return image_cost
input_cost_per_token = _to_float(model_info.get("input_cost_per_token"))
output_cost_per_token = _to_float(model_info.get("output_cost_per_token"))
input_tokens = _estimate_input_tokens(
@ -858,6 +865,40 @@ def _estimate_request_max_cost_for_model(
return cost
def _estimate_image_generation_cost(
request_body: dict,
model_info: Dict[str, Any],
) -> Optional[float]:
"""
Reserve `n × per-image cost` for image-generation requests so concurrent
requests against a depleted budget cannot all slip past the admission gate
onto the provider. Token-based pricing (e.g. gpt-image-1) is handled by
the chat-route token path; per-pixel and size/quality-tiered pricing
(DALL-E 2 size variants, premium tiers) are not handled here and fall
through to read-time enforcement.
The "output" vs "input" cost-per-image naming is inconsistent across
providers — OpenAI's dall-e-3 entry uses ``input_cost_per_image`` while
aiml/dall-e-3 uses ``output_cost_per_image`` — so both are summed.
"""
output_cost_per_image = _to_float(model_info.get("output_cost_per_image"))
input_cost_per_image = _to_float(model_info.get("input_cost_per_image"))
is_image_gen = (
model_info.get("mode") == "image_generation"
or output_cost_per_image is not None
or input_cost_per_image is not None
)
if not is_image_gen:
return None
cost_per_image = (output_cost_per_image or 0.0) + (input_cost_per_image or 0.0)
if cost_per_image <= 0:
return None
n = _to_int(request_body.get("n")) or 1
return cost_per_image * max(n, 1)
def _get_model_cost_info(
model: str,
llm_router: Optional[Router],

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

@ -699,6 +699,154 @@ async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequest
await release_budget_reservation(reservation)
@pytest.mark.asyncio
async def test_should_reserve_image_generation_cost_per_image(
spend_counter_state,
):
"""Image-generation requests reserve `n × per-image cost` so concurrent
requests against a depleted budget cannot all bypass the admission gate.
The OpenAI ``dall-e-3`` entry exposes the per-image price as
``input_cost_per_image`` (a naming quirk), while other providers use
``output_cost_per_image`` — both must be honored."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-image-gen",
spend=0.0,
max_budget=10.0,
)
await key_cache.async_set_cache(key="key-image-gen", value=valid_token)
request_body = {"model": "dall-e-3", "prompt": "a cat", "n": 3}
with patch(
"litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info",
return_value={
"mode": "image_generation",
"input_cost_per_image": 0.04,
},
):
reservation = await reserve_budget_for_request(
request_body=request_body,
route="/v1/images/generations",
llm_router=None,
valid_token=valid_token,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert reservation is not None
assert reservation["reserved_cost"] == pytest.approx(0.12) # 3 × $0.04
await release_budget_reservation(reservation)
@pytest.mark.asyncio
async def test_should_reject_concurrent_image_request_against_depleted_budget(
spend_counter_state,
):
"""Greptile P1 regression: with image-gen reservation in place, a second
concurrent image request against a budget already pinned at the cap by
the first reservation must raise BudgetExceededError instead of
silently reaching the provider."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-image-deplete",
spend=0.0,
team_id="team-image-deplete",
)
team_object = LiteLLM_TeamTable(
team_id="team-image-deplete",
max_budget=0.04,
spend=0.0,
)
await key_cache.async_set_cache(
key=f"team_id:{team_object.team_id}",
value=team_object,
)
request_body = {"model": "dall-e-3", "prompt": "a cat"}
with patch(
"litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info",
return_value={
"mode": "image_generation",
"input_cost_per_image": 0.04,
},
):
first = await reserve_budget_for_request(
request_body=request_body,
route="/v1/images/generations",
llm_router=None,
valid_token=valid_token,
team_object=team_object,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert first is not None
with pytest.raises(litellm.BudgetExceededError):
await reserve_budget_for_request(
request_body=request_body,
route="/v1/images/generations",
llm_router=None,
valid_token=valid_token,
team_object=team_object,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await release_budget_reservation(first)
@pytest.mark.asyncio
async def test_should_skip_reservation_for_per_pixel_image_model(
spend_counter_state,
):
"""DALL-E 2-style per-pixel pricing depends on the requested ``size``,
which we don't decode here. Fall through to read-time enforcement
rather than guess."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
valid_token = UserAPIKeyAuth(
token="key-image-per-pixel",
spend=0.0,
max_budget=1.0,
)
await key_cache.async_set_cache(key="key-image-per-pixel", value=valid_token)
request_body = {"model": "dall-e-2", "prompt": "a cat", "size": "256x256"}
with patch(
"litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info",
return_value={
"mode": "image_generation",
"input_cost_per_pixel": 2.4414e-07,
"output_cost_per_pixel": 0.0,
},
):
reservation = await reserve_budget_for_request(
request_body=request_body,
route="/v1/images/generations",
llm_router=None,
valid_token=valid_token,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert reservation is None
def test_should_start_window_without_reset_at_at_duration_boundary():
before = datetime.now(timezone.utc) - timedelta(hours=1)