From 1ceae6fc43a28a3aca12322ea7a1b05bcb55832c Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Thu, 24 Sep 2026 00:19:45 +0000 Subject: [PATCH] feat(cost): bill azure_ai image edits by measured reference megapixels --- .../image_generation/cost_calculator.py | 105 ++++++----- .../llms/vertex_ai/videos/transformation.py | 6 +- litellm/types/router.py | 1 + litellm/utils.py | 2 + .../test_azure_ai_flux2_image_generation.py | 13 +- .../image_generation/test_cost_calculator.py | 163 ++++++++++++++++++ tests/test_litellm/types/test_router.py | 24 ++- 7 files changed, 259 insertions(+), 55 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 08b732197f3..a0ddbd85618 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -3,22 +3,45 @@ from typing import Any, Final import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import ( - _get_cost_per_unit, + _get_cost_per_unit, # pyright: ignore[reportPrivateUsage] # shared rate-resolution helper used by sibling calculators calculate_image_response_cost_from_usage, resolve_image_model_info, ) from litellm.types.utils import ImageResponse, ModelInfo -def _input_cost_per_pixel(resolved: ModelInfo) -> float: - deployment_price: Final = _get_cost_per_unit(resolved, "input_cost_per_pixel", default_value=None) - if deployment_price is not None: - return deployment_price - model_cost_key: Final = resolved.get("key") - shared_entry: Final = litellm.model_cost.get(model_cost_key) if model_cost_key is not None else None - if shared_entry is None: - return 0.0 - return shared_entry.get("input_cost_per_pixel") or 0.0 +def _rate(table: ModelInfo, key: str) -> float | None: + return _get_cost_per_unit(table, key, default_value=None) + + +def _size_pixels(size: str | None) -> int: + if size is None: + return 0 + for separator in ("x", "-x-"): + if separator in size: + parts = size.split(separator) + if len(parts) != 2: + continue + try: + width = int(parts[0]) + height = int(parts[1]) + except ValueError: + continue + if width > 0 and height > 0: + return width * height + continue + return 0 + + +def _generated_pixels( + optional_params: Mapping[str, object] | None, size: str | None, image_response: ImageResponse +) -> int: + width: Final = optional_params.get("width") if optional_params else None + height: Final = optional_params.get("height") if optional_params else None + if type(width) is int and type(height) is int and width > 0 and height > 0: + return width * height + raw_size: Final = size or image_response.size + return _size_pixels(raw_size if isinstance(raw_size, str) else None) def cost_calculator( @@ -32,44 +55,36 @@ def cost_calculator( """ Azure AI image generation cost calculator """ - _model_info: Final = resolve_image_model_info( + if not isinstance(image_response, ImageResponse): + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + + resolved: Final = resolve_image_model_info( model=model, custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, model_info=model_info, ) - if isinstance(image_response, ImageResponse): - token_based_cost: Final = calculate_image_response_cost_from_usage( - model=model, - image_response=image_response, - custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, - model_info=_model_info, - ) - if token_based_cost is not None: - return token_based_cost + usage_cost: Final = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, + model_info=resolved, + ) + if usage_cost is not None: + return usage_cost - num_images: Final = n if n is not None else len(image_response.data or ()) - output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0 - if output_cost_per_image: - return output_cost_per_image * num_images - - if _input_cost_per_pixel(_model_info): - from litellm.cost_calculator import default_image_cost_calculator - - width: Final = optional_params.get("width") if optional_params else None - height: Final = optional_params.get("height") if optional_params else None - pixel_size: Final = ( - f"{width}x{height}" - if type(width) is int and type(height) is int and width > 0 and height > 0 - else size or image_response.size - ) - return default_image_cost_calculator( - model=_model_info.get("key", model), - custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, - size=pixel_size, - n=num_images, - model_info=model_info, - ) - return 0.0 - - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + data: Final = image_response.data + num_images: Final = n if n is not None else (len(data) if isinstance(data, list) else 0) + generated_meters: Final = ( + ("output_cost_per_image", num_images), + ("input_cost_per_pixel", _generated_pixels(optional_params, size, image_response) * num_images), + ) + generated: Final = next( + (rate * units for key, units in generated_meters if (rate := _rate(resolved, key)) is not None), + 0.0, + ) + reference_pixels: Final = getattr(image_response, "_reference_pixels", None) + reference: Final = (_rate(resolved, "input_cost_per_reference_pixel") or 0.0) * ( + reference_pixels if isinstance(reference_pixels, int) and not isinstance(reference_pixels, bool) else 0 + ) + return generated + reference diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 6c29059f0a9..4ee67cb2618 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -257,9 +257,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): @staticmethod def _supports_resolution_inference(model: str) -> bool: - model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" - model_info: Final = litellm.model_cost.get(model_key) - return model_info is not None and model_info.get("output_cost_per_second_1080p") is not None + model_name: Final = model.removeprefix("vertex_ai/") + model_info: Final = litellm.get_model_info(model=model_name, custom_llm_provider="vertex_ai") + return model_info.get("output_cost_per_second_1080p") is not None def validate_environment( self, diff --git a/litellm/types/router.py b/litellm/types/router.py index 57bd4263894..62fc4925a27 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -791,6 +791,7 @@ class ModelGroupInfo(BaseModel): input_cost_per_token: float | None = None output_cost_per_token: float | None = None input_cost_per_pixel: float | None = None + input_cost_per_reference_pixel: float | None = None mode: ( str | Literal["chat", "embedding", "completion", "image_generation", "audio_transcription", "rerank", "moderations"] diff --git a/litellm/utils.py b/litellm/utils.py index f3b9070ecff..1dd4aa48598 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6201,6 +6201,8 @@ def _get_model_info_helper( output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), + input_cost_per_pixel=_model_info.get("input_cost_per_pixel", None), + input_cost_per_reference_pixel=_model_info.get("input_cost_per_reference_pixel", None), output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index fdd21c87732..409364b8af4 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -127,7 +127,12 @@ def test_flux2_flex_model_info(): assert model_info["max_input_tokens"] == 32000 assert model_info["max_tokens"] == 32000 assert model_info["supported_endpoints"] == ["/v1/images/generations", "/v1/images/edits"] - assert catalog_info["input_cost_per_pixel"] == 5e-08 + assert catalog_info["input_cost_per_reference_pixel"] == catalog_info["input_cost_per_pixel"] + assert catalog_info["input_cost_per_pixel"] * 1024 * 1024 == pytest.approx(0.05), ( + "Azure Retail Prices API, product 'Azure BFL Flux Models', meters 'Flex Megapixel' and " + "'Flex Ref Megapixel' are $0.05 per MP where 1 MP = 1024x1024 pixels; confirmed against " + "Azure Cost Management usage on 2026-09-22 (1024x1024 image metered as 1.0 MP)" + ) assert catalog_info["supported_modalities"] == ["text", "image"] assert catalog_info["supported_output_modalities"] == ["image"] @@ -148,7 +153,7 @@ def test_flux2_flex_cost_uses_generated_megapixels(): call_type="image_generation", ) - assert cost == pytest.approx(5e-08 * 2048 * 1024 * 2) + assert cost == pytest.approx(litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] * 2048 * 1024 * 2) @pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro")) @@ -189,7 +194,7 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi completion_response=response, optional_params=params, call_type="image_generation", - ) == pytest.approx(5e-08 * 2048 * 1024 * 2) + ) == pytest.approx(litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] * 2048 * 1024 * 2) def test_flux2_flex_cost_accepts_lowercase_model_spelling(): @@ -202,7 +207,7 @@ def test_flux2_flex_cost_accepts_lowercase_model_spelling(): call_type="image_generation", ) - assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2) + assert cost == pytest.approx(litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] * 1536 * 1024 * 2) def test_flux2_flex_cost_prefers_deployment_input_cost_per_pixel() -> None: diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py b/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py new file mode 100644 index 00000000000..7f79e518bf3 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_cost_calculator.py @@ -0,0 +1,163 @@ +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.llms.azure_ai.image_generation.cost_calculator import cost_calculator +from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails +from litellm.utils import _invalidate_model_cost_lowercase_map + +REFERENCE_PIXELS: Final = 2 * 1024 * 1024 + + +@pytest.fixture(autouse=True) +def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + yield + litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() + + +def _edit_response(reference_pixels: int | None = REFERENCE_PIXELS, **kwargs: object) -> ImageResponse: + response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n")], size="1024x1024", **kwargs) + response._reference_pixels = reference_pixels + return response + + +def _edit_cost(response: ImageResponse, **kwargs: object) -> float: + return CostCalculatorUtils.route_image_generation_cost_calculator( + model="FLUX.2-flex", + completion_response=response, + custom_llm_provider="azure_ai", + size="1024x1024", + call_type="image_edit", + **kwargs, + ) + + +def test_get_model_info_surfaces_flux2_flex_pixel_rates() -> None: + model_info = litellm.get_model_info(model="FLUX.2-flex", custom_llm_provider="azure_ai") + catalog_info = litellm.model_cost["azure_ai/FLUX.2-flex"] + + assert model_info["input_cost_per_pixel"] == catalog_info["input_cost_per_pixel"] + assert model_info["input_cost_per_reference_pixel"] == catalog_info["input_cost_per_reference_pixel"] + + +def test_flux2_flex_catalog_pixel_rates_match_reference_rates() -> None: + catalog_info = litellm.model_cost["azure_ai/FLUX.2-flex"] + + assert catalog_info["input_cost_per_reference_pixel"] == catalog_info["input_cost_per_pixel"] + assert catalog_info["input_cost_per_pixel"] * 1024 * 1024 == pytest.approx(0.05), ( + "Azure Retail Prices API, product 'Azure BFL Flux Models', meters 'Flex Megapixel' and " + "'Flex Ref Megapixel' are $0.05 per MP where 1 MP = 1024x1024 pixels; confirmed against " + "Azure Cost Management usage on 2026-09-22 (1024x1024 image metered as 1.0 MP)" + ) + + +def test_edit_cost_adds_reference_pixels_to_generated_pixels() -> None: + catalog_rate: Final = litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] + + cost: Final = _edit_cost(_edit_response()) + + assert cost == pytest.approx(catalog_rate * 1024 * 1024 + catalog_rate * REFERENCE_PIXELS) + + +def test_edit_cost_prefers_deployment_rates_over_catalog() -> None: + cost: Final = _edit_cost( + _edit_response(), model_info={"input_cost_per_pixel": 2e-07, "input_cost_per_reference_pixel": 3e-07} + ) + + assert cost == pytest.approx(2e-07 * 1024 * 1024 + 3e-07 * REFERENCE_PIXELS) + + +def test_edit_cost_honors_explicit_zero_reference_rate() -> None: + cost: Final = _edit_cost( + _edit_response(), model_info={"input_cost_per_pixel": 2e-07, "input_cost_per_reference_pixel": 0.0} + ) + + assert cost == pytest.approx(2e-07 * 1024 * 1024) + + +def test_edit_cost_honors_explicit_zero_generated_rate() -> None: + cost: Final = _edit_cost( + _edit_response(), model_info={"input_cost_per_pixel": 0.0, "input_cost_per_reference_pixel": 1e-07} + ) + + assert cost == pytest.approx(1e-07 * REFERENCE_PIXELS) + + +def test_per_image_rate_beats_per_pixel_rate(monkeypatch: pytest.MonkeyPatch) -> None: + cost: Final = _edit_cost( + _edit_response(), + model_info={ + "output_cost_per_image": 0.04, + "input_cost_per_pixel": 1e-07, + "input_cost_per_reference_pixel": 1.5e-08, + }, + ) + + assert cost == pytest.approx(0.04 + 1.5e-08 * REFERENCE_PIXELS) + + +def test_usage_short_circuits_pixel_billing() -> None: + response: Final = _edit_response( + usage=ImageUsage( + input_tokens=150, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=100, text_tokens=50), + output_tokens=1000, + total_tokens=1150, + ) + ) + + cost: Final = _edit_cost( + response, + model_info={ + "input_cost_per_pixel": 5e-08, + "input_cost_per_reference_pixel": 5e-08, + "input_cost_per_token": 1e-05, + "input_cost_per_image_token": 2e-05, + "output_cost_per_image_token": 4e-05, + "output_cost_per_token": 4e-05, + }, + ) + + assert cost == pytest.approx(50 * 1e-05 + 100 * 2e-05 + 1000 * 4e-05) + + +def test_edit_without_measurement_bills_generated_pixels_only() -> None: + catalog_rate: Final = litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] + + cost: Final = _edit_cost(_edit_response(reference_pixels=None)) + + assert cost == pytest.approx(catalog_rate * 1024 * 1024) + + +def test_optional_params_dimensions_beat_response_size() -> None: + cost: Final = _edit_cost( + _edit_response(reference_pixels=None), + optional_params={"width": 2048, "height": 1024}, + ) + + assert cost == pytest.approx(litellm.model_cost["azure_ai/FLUX.2-flex"]["input_cost_per_pixel"] * 2048 * 1024) + + +def test_unlisted_model_bills_deployment_rates() -> None: + cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator( + model="unlisted-flux-deployment", + completion_response=_edit_response(), + custom_llm_provider="azure_ai", + size="1024x1024", + call_type="image_edit", + model_info={"input_cost_per_pixel": 1e-07, "input_cost_per_reference_pixel": 1e-07}, + ) + + assert cost == pytest.approx(1e-07 * 1024 * 1024 + 1e-07 * REFERENCE_PIXELS) + + +def test_non_image_response_raises() -> None: + with pytest.raises(ValueError): + cost_calculator(model="FLUX.2-flex", image_response=object()) diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index 4d4c326d1ca..ffca7c91f67 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -229,7 +229,25 @@ def test_model_info_rejects_offset_aware_access_window_times(): with pytest.raises(ValidationError): ModelInfo( id="x", - access_windows=[ - {"start": "22:00+05:00", "end": "06:00", "timezone": "UTC", "team_ids": ["t"]} - ], + access_windows=[{"start": "22:00+05:00", "end": "06:00", "timezone": "UTC", "team_ids": ["t"]}], ) + + +def test_model_group_info_surfaces_input_cost_per_reference_pixel(): + from typing import Final + + from litellm import Router + + configured_rate: Final = 3e-7 + router: Final = Router( + model_list=[ + { + "model_name": "flux-edit", + "litellm_params": {"model": "azure_ai/FLUX.2-flex"}, + "model_info": {"id": "dep-flux-1", "input_cost_per_reference_pixel": configured_rate}, + } + ] + ) + group: Final = router.get_model_group_info("flux-edit") + assert group is not None + assert group.input_cost_per_reference_pixel == configured_rate