fix(accounting): retain billable image dimensions for replay

This commit is contained in:
Joshua Valluru 2026-10-31 15:41:08 -07:00
parent 8b73a8afe7
commit 0010c2d443
2 changed files with 42 additions and 16 deletions

View file

@ -1155,12 +1155,23 @@ def _realtime_pricing_models(value: Mapping[str, JsonValue]) -> Iterator[str]:
pending.extend(nested for key in _REALTIME_CHILDREN if isinstance(nested := fields.get(key), dict))
def _image_cost_dimensions(image: object) -> Mapping[str, object]:
fields: Final = _fields(_response_fields(image).get("provider_specific_fields"))
return {
"provider_specific_fields": {
name: value
for name, value in fields.items()
if name in {"width", "height"} and isinstance(value, int) and not isinstance(value, bool)
}
}
def _pricing_inputs(kwargs: Mapping[str, object], logging_obj: object, response: object) -> Mapping[str, object]:
from litellm.litellm_core_utils.litellm_logging import Logging
return {
"image_response": {
"data": tuple(dict[str, JsonValue]() for _ in response.data or ()),
"data": tuple(_image_cost_dimensions(image) for image in response.data or ()),
"usage": _response_fields(response.usage) if response.usage is not None else None,
"size": response.size,
"quality": response.quality,

View file

@ -1609,9 +1609,12 @@ async def test_live_adjustment_reload_is_captured_and_replay_remains_frozen(monk
@pytest.mark.asyncio
@pytest.mark.parametrize("token_pricing", (True, False))
@pytest.mark.parametrize(
("pricing", "width", "expected"),
(("tokens", 2048, 0.2), ("images", 2048, 0.5), ("pixels", 2048, 0.4), ("pixels", 2049, 0.6)),
)
async def test_image_capture_replays_the_existing_image_calculator(
monkeypatch: pytest.MonkeyPatch, token_pricing: bool
monkeypatch: pytest.MonkeyPatch, pricing: str, width: int, expected: float
) -> None:
import litellm
from litellm.proxy.spend_tracking.usage_capture import capture_logging_usage
@ -1622,30 +1625,42 @@ async def test_image_capture_replays_the_existing_image_calculator(
monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.1})
monkeypatch.setattr(litellm, "cost_margin_config", {"global": {"fixed_amount": 0.2}})
model: Final = "gpt-image-capture" if token_pricing else "dall-e-capture"
model: Final = {"tokens": "gpt-image-capture", "images": "dall-e-capture", "pixels": "fal-ai/flux/dev"}[pricing]
provider: Final = "fal_ai" if pricing == "pixels" else "openai"
catalog_key: Final = f"{provider}/{model}" if pricing == "pixels" else model
monkeypatch.setitem(
litellm.model_cost,
model,
catalog_key,
{
"litellm_provider": "openai",
"litellm_provider": provider,
"mode": "image_generation",
**(
{"input_cost_per_token": 0.01, "output_cost_per_image_token": 0.02}
if token_pricing
else {"output_cost_per_image": 0.25}
),
**{
"tokens": {"input_cost_per_token": 0.01, "output_cost_per_image_token": 0.02},
"images": {"output_cost_per_image": 0.25},
"pixels": {"output_cost_per_image": 0.1, "output_cost_per_pixel": 0.1 / (1024 * 1024)},
}[pricing],
},
)
logger: Final = _logger("aimage_generation")
logger.model = model
logger.update_environment_variables(
model=model, custom_llm_provider="openai", litellm_params={}, optional_params={}
model=model, custom_llm_provider=provider, litellm_params={}, optional_params={}
)
response: Final = ImageResponse.model_validate(
{
"created": 1,
"model": model,
"data": [{"b64_json": "private-image-content"}] * 2,
"data": [
{
"b64_json": "private-image-content",
"provider_specific_fields": {
"width": width,
"height": 1024,
"secret": "private-image-content",
},
}
]
* 2,
"size": "1024x1024",
"quality": "standard",
"usage": {
@ -1654,12 +1669,12 @@ async def test_image_capture_replays_the_existing_image_calculator(
"total_tokens": 15,
"input_tokens_details": {"text_tokens": 10, "image_tokens": 0},
}
if token_pricing
if pricing == "tokens"
else None,
}
)
source: Final = logger.response_cost_calculator(response, ignore_stamped_cost=True)
assert source == pytest.approx(0.2 if token_pricing else 0.5)
assert source == pytest.approx(expected)
database: Final = StoredUsage()
meter: Final = DurableUsageMeter(database.insert, UsageCaptureSettings(), wait_for_batch=_manual_flush)
try:
@ -1667,7 +1682,7 @@ async def test_image_capture_replays_the_existing_image_calculator(
await meter.close()
envelope: Final = database.rows[-1].envelope
assert "private-image-content" not in envelope
monkeypatch.delitem(litellm.model_cost, model)
monkeypatch.delitem(litellm.model_cost, catalog_key)
replay: Final = price_usage(envelope)
assert isinstance(replay, PricingSnapshot)
assert float(replay.total_cost) == pytest.approx(source)