From 53adf6575f6adfc101aa2da7ea62e423c66dfc7e Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 27 Oct 2026 13:16:35 -0700 Subject: [PATCH 1/4] fix(azure_ai): bill FLUX.2 images from Azure's reported megapixels Azure meters FLUX.2 by megapixel, with the first output megapixel priced apart from the rest and reference images billed per megapixel, all rounded up. Every FLUX.2 response reports the usage in request_meta (output_mp, input_mp). Record that usage on generation, edit and passthrough responses and multiply it by new catalog prices, so logged spend follows what Azure bills. Responses without request_meta keep today's per-image price, and a deployment that sets only its own per-image price keeps it. --- .../crates/model-catalog/src/model_info.rs | 6 + litellm/llms/azure/azure.py | 14 +- .../image_edit/flux2_transformation.py | 14 ++ .../image_generation/cost_calculator.py | 60 ++++++ .../azure_ai/passthrough/transformation.py | 10 +- ...odel_prices_and_context_window_backup.json | 7 +- litellm/types/utils.py | 6 + litellm/utils.py | 3 + model_prices_and_context_window.json | 7 +- ...test_azure_ai_image_edit_transformation.py | 25 +++ .../test_azure_ai_flux2_image_generation.py | 184 +++++++++++++++++- ...est_azure_ai_passthrough_transformation.py | 37 ++++ tests/unit/test_utils.py | 6 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 ++ 14 files changed, 384 insertions(+), 7 deletions(-) diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index ebee2c64140..38737fb2569 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -196,6 +196,8 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_image_token_batches: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub input_cost_per_megapixel: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_pixel: Option, #[serde(skip_serializing_if = "Option::is_none")] pub input_cost_per_query: Option, @@ -314,6 +316,8 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_character_above_128k_tokens: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_first_megapixel: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image: Option, #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_1024: Option, @@ -344,6 +348,8 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_image_token: Option, #[serde(skip_serializing_if = "Option::is_none")] + pub output_cost_per_megapixel: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub output_cost_per_pixel: Option, /// USD per reasoning/thinking token, when billed separately. #[serde(skip_serializing_if = "Option::is_none")] diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 02f1ba3472a..f6a1bf251c9 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1195,7 +1195,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) if provider_config is not None: - return provider_config.transform_image_generation_response( + from litellm.llms.azure_ai.image_generation.cost_calculator import record_megapixel_usage + + image_response: Final = provider_config.transform_image_generation_response( model=data.get("model", "dall-e-2"), raw_response=httpx_response, model_response=model_response or ImageResponse(), @@ -1205,6 +1207,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=data, encoding=litellm.encoding, ) + record_megapixel_usage(image_response, httpx_response.content) + return image_response else: response = httpx_response.json() @@ -1351,12 +1355,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={"complete_input_dict": data}, original_response=response, ) - # return response - return convert_to_model_response_object( + from litellm.llms.azure_ai.image_generation.cost_calculator import record_megapixel_usage + + converted_response: Final = convert_to_model_response_object( response_object=response, model_response_object=model_response, response_type="image_generation", ) + record_megapixel_usage(converted_response, httpx_response.content) except AzureOpenAIError as e: raise e except Exception as e: @@ -1365,6 +1371,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): raise AzureOpenAIError(status_code=error_code, message=str(e)) else: raise AzureOpenAIError(status_code=500, message=str(e)) + else: + return converted_response def audio_speech( self, diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 9bd476efdff..3141de38384 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -4,6 +4,7 @@ from io import BufferedReader from types import MappingProxyType from typing import Any, Final +import httpx from httpx._types import RequestFiles import litellm @@ -11,14 +12,17 @@ from litellm.llms.azure_ai.common_utils import ( AzureFoundryModelInfo, get_azure_ai_auth_headers, ) +from litellm.llms.azure_ai.image_generation.cost_calculator import record_megapixel_usage from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) +from litellm.llms.base_llm.image_edit.transformation import LiteLLMLoggingObj from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.openai import FileTypes from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): @@ -135,6 +139,16 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): return base64.b64encode(image_bytes).decode("utf-8") + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ImageResponse: + image_response: Final = super().transform_image_edit_response(model, raw_response, logging_obj) + record_megapixel_usage(image_response, raw_response.content) + return image_response + def get_complete_url( self, model: str, diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index beed3881bdf..cfb564d7027 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,7 +1,11 @@ +import math from collections.abc import Mapping from typing import Final +from pydantic import BaseModel, ConfigDict, Field, ValidationError + import litellm +from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_image_response_cost_from_usage, get_cost_per_unit, @@ -9,6 +13,55 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( ) from litellm.types.utils import ImageResponse, ModelInfo +_MEGAPIXEL_USAGE_HIDDEN_PARAM: Final = "megapixel_usage" +_MEGAPIXEL_PRICE_KEYS: Final = ( + "output_cost_per_first_megapixel", + "output_cost_per_megapixel", + "input_cost_per_megapixel", +) + + +class _MegapixelUsage(BaseModel): + model_config = ConfigDict(frozen=True, allow_inf_nan=False) + + input_mp: float | None = Field(default=None, ge=0) + output_mp: float = Field(gt=0) + + +class _MeteredImageResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + request_meta: _MegapixelUsage + + +def record_megapixel_usage(response: object, response_body: bytes | Mapping[str, object]) -> None: + try: + usage: Final = ( + _MeteredImageResponse.model_validate_json(response_body) + if isinstance(response_body, bytes) + else _MeteredImageResponse.model_validate(response_body) + ).request_meta + except ValidationError: + return + set_hidden_param(response, _MEGAPIXEL_USAGE_HIDDEN_PARAM, usage.model_dump()) + + +def _megapixel_cost(prices: ModelInfo, image_response: ImageResponse) -> float | None: + output_rate: Final = get_cost_per_unit(prices, "output_cost_per_megapixel", default_value=None) + if output_rate is None: + return None + try: + usage: Final = _MegapixelUsage.model_validate(image_response.hidden_params.get(_MEGAPIXEL_USAGE_HIDDEN_PARAM)) + except ValidationError: + return None + first_rate: Final = get_cost_per_unit(prices, "output_cost_per_first_megapixel", default_value=None) + input_rate: Final = get_cost_per_unit(prices, "input_cost_per_megapixel") or 0.0 + return ( + (output_rate if first_rate is None else first_rate) + + output_rate * (math.ceil(usage.output_mp) - 1) + + input_rate * math.ceil(usage.input_mp or 0) + ) + def _input_cost_per_pixel(resolved: ModelInfo) -> float: deployment_price: Final = get_cost_per_unit(resolved, "input_cost_per_pixel", default_value=None) @@ -48,6 +101,13 @@ def cost_calculator( if token_based_cost is not None: return token_based_cost + bills_megapixels: Final = model_info is None or any( + model_info.get(key) is not None for key in _MEGAPIXEL_PRICE_KEYS + ) + megapixel_cost: Final = _megapixel_cost(_model_info, image_response) if bills_megapixels else None + if megapixel_cost is not None: + return megapixel_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: diff --git a/litellm/llms/azure_ai/passthrough/transformation.py b/litellm/llms/azure_ai/passthrough/transformation.py index 35b7642d7de..2c1af848167 100644 --- a/litellm/llms/azure_ai/passthrough/transformation.py +++ b/litellm/llms/azure_ai/passthrough/transformation.py @@ -13,6 +13,7 @@ from litellm.llms.azure_ai.common_utils import ( api_key_header_for_base, get_azure_ai_auth_headers, ) +from litellm.llms.azure_ai.image_generation.cost_calculator import record_megapixel_usage from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.passthrough.transformation import ( BasePassthroughConfig, @@ -85,9 +86,16 @@ def relay_query_params( return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version}) +def metered_image_response(body: Mapping[str, object]) -> ImageResponse: + image_response: Final = ImageResponse.model_validate(body) + record_megapixel_usage(image_response, body) + return image_response + + FOUNDRY_RELAY_SHAPES: Final = ( RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate), - RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate), + RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, metered_image_response), + RelayShape("/providers/blackforestlabs/v1/flux-2-flex", CallTypes.aimage_generation, metered_image_response), ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7cd4eadca6e..87b6d1bc902 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11994,20 +11994,25 @@ ] }, "azure_ai/flux.2-pro": { + "input_cost_per_megapixel": 0.015, "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", + "output_cost_per_first_megapixel": 0.03, "output_cost_per_image": 0.04, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "output_cost_per_megapixel": 0.015, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", "supported_endpoints": [ "/v1/images/generations" ] }, "azure_ai/FLUX.2-flex": { + "input_cost_per_megapixel": 0.05, "input_cost_per_pixel": 5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", + "output_cost_per_megapixel": 0.05, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", "supported_endpoints": [ "/v1/images/generations", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 7e8645fd0dd..947a7a3ec37 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -428,6 +428,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_character_above_128k_tokens: float | None # only for vertex ai models output_cost_per_image: float | None output_cost_per_pixel: ReadOnly[float | None] + output_cost_per_first_megapixel: ReadOnly[float | None] + output_cost_per_megapixel: ReadOnly[float | None] + input_cost_per_megapixel: ReadOnly[float | None] output_cost_per_image_token: float | None output_cost_per_image_token_batches: ReadOnly[float | None] output_cost_per_video_token: float | None # for gemini omni models with video output @@ -4001,6 +4004,9 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_image_1536: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None + output_cost_per_first_megapixel: float | None = None + output_cost_per_megapixel: float | None = None + input_cost_per_megapixel: float | None = None # Include all ModelInfoBase fields as optional # This allows any model_info parameter to be set in litellm_params diff --git a/litellm/utils.py b/litellm/utils.py index 6f0b1ca2e3e..0226706d28b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6954,6 +6954,9 @@ def get_model_info_helper( 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), output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None), + output_cost_per_first_megapixel=_model_info.get("output_cost_per_first_megapixel", None), + output_cost_per_megapixel=_model_info.get("output_cost_per_megapixel", None), + input_cost_per_megapixel=_model_info.get("input_cost_per_megapixel", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_cost_per_image_token_batches=_model_info.get("output_cost_per_image_token_batches", None), output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7cd4eadca6e..87b6d1bc902 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11994,20 +11994,25 @@ ] }, "azure_ai/flux.2-pro": { + "input_cost_per_megapixel": 0.015, "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", + "output_cost_per_first_megapixel": 0.03, "output_cost_per_image": 0.04, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "output_cost_per_megapixel": 0.015, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", "supported_endpoints": [ "/v1/images/generations" ] }, "azure_ai/FLUX.2-flex": { + "input_cost_per_megapixel": 0.05, "input_cost_per_pixel": 5e-08, "litellm_provider": "azure_ai", "max_input_tokens": 32000, "mode": "image_generation", + "output_cost_per_megapixel": 0.05, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/", "supported_endpoints": [ "/v1/images/generations", diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 51ba2c34cd7..8873a963288 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -178,6 +178,31 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ assert response._hidden_params["response_cost"] == pytest.approx(5e-08 * 2048 * 1024 * 2) +@pytest.mark.usefixtures("local_model_cost_map") +def test_flux2_image_edit_bills_the_reported_reference_megapixels() -> None: + metered_body: Final = { + "data": [{"b64_json": "aW1n"}], + "request_meta": {"cost": 999.0, "input_mp": 2.0, "output_mp": 1.0, "total_pixels": 1}, + } + client: Final = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(lambda _: httpx.Response(200, json=metered_body))) + ) + + response: Final = litellm.image_edit( + model="azure_ai/flux.2-pro", + image=[b"first", b"second"], + prompt="Combine them", + api_key="test-key", + api_base="https://example.services.ai.azure.com", + client=client, + ) + + prices: Final = litellm.model_cost["azure_ai/flux.2-pro"] + assert response._hidden_params["response_cost"] == pytest.approx( + prices["output_cost_per_first_megapixel"] + 2 * prices["input_cost_per_megapixel"] + ) + + def test_flux2_image_edit_accepts_and_drops_openai_only_parameters(): optional_params: Final = ImageEditRequestUtils.get_optional_params_image_edit( model="FLUX.2-pro", diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index fdd21c87732..dda98c1b07e 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -1,9 +1,11 @@ -from collections.abc import Mapping +import json +from collections.abc import Callable, Mapping from typing import Final from unittest.mock import MagicMock import httpx import pytest +import respx import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -251,3 +253,183 @@ def test_flux2_response_preserves_mapped_dimensions(): encoding=None, ) assert response.size == "2048x1024" + + +FOUNDRY_BASE: Final = "https://flux-billing.services.ai.azure.com" + + +def _flux2_route(provider_path: str) -> str: + return rf"https://flux-billing\.services\.ai\.azure\.com/providers/blackforestlabs/v1/{provider_path}\?.*" + + +def _metered_body(output_mp: float, input_mp: float) -> Mapping[str, object]: + return { + "created": 1, + "data": [{"b64_json": "aW1n"}], + "request_meta": {"cost": 999.0, "input_mp": input_mp, "output_mp": output_mp, "total_pixels": 1}, + } + + +def _price(model: str, key: str) -> float: + price: Final = litellm.model_cost[model][key] + assert isinstance(price, float) and price > 0 + return price + + +# request_meta pairs Azure returned for FLUX.2 [pro] on 2026-10-27, whose credit charges matched whole megapixels +# rounded up: one 512x512 reference reported input_mp 0.25 and two of them reported 2.0 +@pytest.mark.parametrize( + ("output_mp", "input_mp", "additional_output_megapixels", "input_megapixels"), + [ + (0.25, 0.0, 0, 0), + (1.5, 0.0, 1, 0), + (2.07, 0.0, 2, 0), + (1.0, 0.25, 0, 1), + (1.0, 1.5, 0, 2), + (1.0, 2.0, 0, 2), + ], +) +@respx.mock +def test_flux2_pro_generation_bills_the_reported_megapixels_rounded_up( + output_mp: float, input_mp: float, additional_output_megapixels: int, input_megapixels: int +) -> None: + respx.post(url__regex=_flux2_route("flux-2-pro")).mock( + return_value=httpx.Response(200, json=_metered_body(output_mp, input_mp)) + ) + + response: Final = litellm.image_generation( + model="azure_ai/flux.2-pro", prompt="A red fox", api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response._hidden_params["response_cost"] == pytest.approx( + _price("azure_ai/flux.2-pro", "output_cost_per_first_megapixel") + + additional_output_megapixels * _price("azure_ai/flux.2-pro", "output_cost_per_megapixel") + + input_megapixels * _price("azure_ai/flux.2-pro", "input_cost_per_megapixel") + ) + + +@pytest.mark.asyncio +async def test_async_flux2_pro_generation_bills_the_reported_megapixels( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.post(url__regex=_flux2_route("flux-2-pro")).mock( + return_value=httpx.Response(200, json=_metered_body(1.5, 0.25)) + ) + + response: Final = await litellm.aimage_generation( + model="azure_ai/flux.2-pro", prompt="A red fox", api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response._hidden_params["response_cost"] == pytest.approx( + _price("azure_ai/flux.2-pro", "output_cost_per_first_megapixel") + + _price("azure_ai/flux.2-pro", "output_cost_per_megapixel") + + _price("azure_ai/flux.2-pro", "input_cost_per_megapixel") + ) + + +@respx.mock +def test_flux2_flex_generation_bills_every_reported_megapixel_at_one_rate() -> None: + respx.post(url__regex=_flux2_route("flux-2-flex")).mock( + return_value=httpx.Response(200, json=_metered_body(1.5, 0.25)) + ) + + response: Final = litellm.image_generation( + model="azure_ai/FLUX.2-flex", prompt="A red fox", api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response._hidden_params["response_cost"] == pytest.approx( + 2 * _price("azure_ai/FLUX.2-flex", "output_cost_per_megapixel") + + _price("azure_ai/FLUX.2-flex", "input_cost_per_megapixel") + ) + + +@pytest.mark.parametrize( + ("deployment_prices", "expected_cost"), + [ + ( + {}, + lambda price: price("output_cost_per_first_megapixel") + + price("output_cost_per_megapixel") + + price("input_cost_per_megapixel"), + ), + ({"output_cost_per_image": 0.05}, lambda _: 0.05), + ( + {"output_cost_per_megapixel": 0.02}, + lambda price: price("output_cost_per_first_megapixel") + 0.02 + price("input_cost_per_megapixel"), + ), + ( + {"output_cost_per_first_megapixel": 0.04, "output_cost_per_megapixel": 0.02, "input_cost_per_megapixel": 0.01}, + lambda _: 0.04 + 0.02 + 0.01, + ), + ], +) +@respx.mock +def test_flux2_router_deployment_prices_layer_over_the_catalog_megapixel_prices( + deployment_prices: Mapping[str, float], expected_cost: Callable[[Callable[[str], float]], float] +) -> None: + respx.post(url__regex=_flux2_route("flux-2-pro")).mock( + return_value=httpx.Response(200, json=_metered_body(1.5, 0.25)) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "flux-2-pro", + "litellm_params": { + "model": "azure_ai/flux.2-pro", + "api_base": FOUNDRY_BASE, + "api_key": "test-key", + **deployment_prices, + }, + } + ] + ) + + response: Final = router.image_generation(model="flux-2-pro", prompt="A red fox") + + assert response._hidden_params["response_cost"] == pytest.approx( + expected_cost(lambda key: _price("azure_ai/flux.2-pro", key)) + ) + + +@pytest.mark.parametrize( + "request_meta", + [ + None, + {"input_mp": 0.0, "output_mp": 0.0}, + {"input_mp": -1.0, "output_mp": 1.0}, + {"input_mp": 0.0, "output_mp": float("inf")}, + {"input_mp": 0.0}, + ], +) +@respx.mock +def test_flux2_pro_generation_without_usable_megapixels_keeps_the_per_image_price( + request_meta: Mapping[str, float] | None, +) -> None: + body: Final = {"created": 1, "data": [{"b64_json": "aW1n"}], "request_meta": request_meta} + respx.post(url__regex=_flux2_route("flux-2-pro")).mock( + return_value=httpx.Response(200, content=json.dumps(body).encode()) + ) + + response: Final = litellm.image_generation( + model="azure_ai/flux.2-pro", prompt="A red fox", api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response._hidden_params["response_cost"] == pytest.approx( + _price("azure_ai/flux.2-pro", "output_cost_per_image") + ) + + +@respx.mock +def test_flux2_pro_generation_bills_no_references_when_input_megapixels_are_null() -> None: + body: Final = {"created": 1, "data": [{"b64_json": "aW1n"}], "request_meta": {"input_mp": None, "output_mp": 1.5}} + respx.post(url__regex=_flux2_route("flux-2-pro")).mock(return_value=httpx.Response(200, json=body)) + + response: Final = litellm.image_generation( + model="azure_ai/flux.2-pro", prompt="A red fox", api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response._hidden_params["response_cost"] == pytest.approx( + _price("azure_ai/flux.2-pro", "output_cost_per_first_megapixel") + + _price("azure_ai/flux.2-pro", "output_cost_per_megapixel") + ) diff --git a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py index 24f21751065..0551915cde5 100644 --- a/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py +++ b/tests/unit/llms/azure_ai/passthrough/test_azure_ai_passthrough_transformation.py @@ -1,5 +1,7 @@ import json +from collections.abc import Callable, Mapping from datetime import datetime +from typing import Final from unittest.mock import MagicMock import httpx @@ -542,6 +544,41 @@ def test_flux_2_relay_through_the_provider_route_is_costed_per_image(): assert logging_obj.response_cost_calculator(result=result) == pytest.approx(per_image) +@pytest.mark.parametrize( + ("model", "native_path", "catalog_key", "expected_cost"), + [ + ( + "FLUX.2-pro", + "providers/blackforestlabs/v1/flux-2-pro", + "azure_ai/flux.2-pro", + lambda prices: prices["output_cost_per_first_megapixel"] + + prices["output_cost_per_megapixel"] + + prices["input_cost_per_megapixel"], + ), + ( + "FLUX.2-flex", + "providers/blackforestlabs/v1/flux-2-flex", + "azure_ai/FLUX.2-flex", + lambda prices: 2 * prices["output_cost_per_megapixel"] + prices["input_cost_per_megapixel"], + ), + ], +) +@pytest.mark.usefixtures("local_model_cost_map") +def test_flux_2_relay_is_costed_from_the_reported_megapixels( + model: str, native_path: str, catalog_key: str, expected_cost: Callable[[Mapping[str, float]], float] +) -> None: + metered_body: Final = { + **IMAGE_BODY, + "request_meta": {"cost": 999.0, "input_mp": 0.25, "output_mp": 1.5, "total_pixels": 1}, + } + result, logging_obj = _relay_logging_result(AzureAIPassthroughConfig(), model, native_path, metered_body) + + assert logging_obj.call_type == "aimage_generation" + assert logging_obj.response_cost_calculator(result=result) == pytest.approx( + expected_cost(litellm.model_cost[catalog_key]) + ) + + def test_rejected_rerank_relay_keeps_the_passthrough_object_and_call_type(): result, logging_obj = _relay_logging_result( AzureAIPassthroughConfig(), diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 38287b25169..81f88a7819f 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -824,6 +824,9 @@ def validate_model_cost_values(model_data, exceptions=None): "output_cost_per_image_4K", "input_cost_per_pixel", "output_cost_per_pixel", + "input_cost_per_megapixel", + "output_cost_per_first_megapixel", + "output_cost_per_megapixel", "cost_per_second", "input_cost_per_second", "input_cost_per_text_message", @@ -1047,6 +1050,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "regional_endpoint_uplift_multiplier": {"type": "number"}, "regional_processing_uplift_multiplier_eu": {"type": "number"}, "regional_processing_uplift_multiplier_us": {"type": "number"}, + "input_cost_per_megapixel": {"type": "number"}, "input_cost_per_pixel": {"type": "number"}, "input_cost_per_query": {"type": "number"}, "input_cost_per_request": {"type": "number"}, @@ -1117,6 +1121,8 @@ _COST_MAP_ROW_SCHEMA: Final = { "output_cost_per_image_token": {"type": "number"}, "output_cost_per_image_token_batches": {"type": "number"}, "output_cost_per_video_token": {"type": "number"}, + "output_cost_per_first_megapixel": {"type": "number"}, + "output_cost_per_megapixel": {"type": "number"}, "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, "output_cost_per_second_480p": {"type": "number"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d03f4b4cffd..f40bb728965 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35398,6 +35398,8 @@ export interface components { input_cost_per_image_token?: number | null; /** Input Cost Per Image Token Batches */ input_cost_per_image_token_batches?: number | null; + /** Input Cost Per Megapixel */ + input_cost_per_megapixel?: number | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; /** Input Cost Per Query */ @@ -35513,6 +35515,8 @@ export interface components { output_cost_per_character?: number | null; /** Output Cost Per Character Above 128K Tokens */ output_cost_per_character_above_128k_tokens?: number | null; + /** Output Cost Per First Megapixel */ + output_cost_per_first_megapixel?: number | null; /** Output Cost Per Image */ output_cost_per_image?: number | null; /** Output Cost Per Image 1024 */ @@ -35525,6 +35529,8 @@ export interface components { output_cost_per_image_token?: number | null; /** Output Cost Per Image Token Batches */ output_cost_per_image_token_batches?: number | null; + /** Output Cost Per Megapixel */ + output_cost_per_megapixel?: number | null; /** Output Cost Per Pixel */ output_cost_per_pixel?: number | null; /** Output Cost Per Reasoning Token */ @@ -50779,6 +50785,8 @@ export interface components { input_cost_per_image_token?: number | null; /** Input Cost Per Image Token Batches */ input_cost_per_image_token_batches?: number | null; + /** Input Cost Per Megapixel */ + input_cost_per_megapixel?: number | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; /** Input Cost Per Query */ @@ -50894,6 +50902,8 @@ export interface components { output_cost_per_character?: number | null; /** Output Cost Per Character Above 128K Tokens */ output_cost_per_character_above_128k_tokens?: number | null; + /** Output Cost Per First Megapixel */ + output_cost_per_first_megapixel?: number | null; /** Output Cost Per Image */ output_cost_per_image?: number | null; /** Output Cost Per Image 1024 */ @@ -50906,6 +50916,8 @@ export interface components { output_cost_per_image_token?: number | null; /** Output Cost Per Image Token Batches */ output_cost_per_image_token_batches?: number | null; + /** Output Cost Per Megapixel */ + output_cost_per_megapixel?: number | null; /** Output Cost Per Pixel */ output_cost_per_pixel?: number | null; /** Output Cost Per Reasoning Token */ From aef764a37216bfd77ad475257dfd1cc498aba1da Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 27 Oct 2026 23:28:48 +0000 Subject: [PATCH 2/4] fix(azure_ai): bill FLUX.2 first megapixel per returned image --- .../image_generation/cost_calculator.py | 10 ++--- .../test_azure_ai_flux2_image_generation.py | 44 ++++++++++++++++--- 2 files changed, 43 insertions(+), 11 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index cfb564d7027..5f421a7e2cd 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -56,11 +56,11 @@ def _megapixel_cost(prices: ModelInfo, image_response: ImageResponse) -> float | return None first_rate: Final = get_cost_per_unit(prices, "output_cost_per_first_megapixel", default_value=None) input_rate: Final = get_cost_per_unit(prices, "input_cost_per_megapixel") or 0.0 - return ( - (output_rate if first_rate is None else first_rate) - + output_rate * (math.ceil(usage.output_mp) - 1) - + input_rate * math.ceil(usage.input_mp or 0) - ) + images: Final = max(len(image_response.data or ()), 1) + output_mp_per_image: Final = usage.output_mp / images + return images * ( + (output_rate if first_rate is None else first_rate) + output_rate * (math.ceil(output_mp_per_image) - 1) + ) + input_rate * math.ceil(usage.input_mp or 0) def _input_cost_per_pixel(resolved: ModelInfo) -> float: diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index dda98c1b07e..5a010f81556 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -262,10 +262,10 @@ def _flux2_route(provider_path: str) -> str: return rf"https://flux-billing\.services\.ai\.azure\.com/providers/blackforestlabs/v1/{provider_path}\?.*" -def _metered_body(output_mp: float, input_mp: float) -> Mapping[str, object]: +def _metered_body(output_mp: float, input_mp: float, num_images: int = 1) -> Mapping[str, object]: return { "created": 1, - "data": [{"b64_json": "aW1n"}], + "data": [{"b64_json": "aW1n"} for _ in range(num_images)], "request_meta": {"cost": 999.0, "input_mp": input_mp, "output_mp": output_mp, "total_pixels": 1}, } @@ -308,6 +308,32 @@ def test_flux2_pro_generation_bills_the_reported_megapixels_rounded_up( ) +@pytest.mark.parametrize( + ("output_mp", "additional_output_megapixels"), + [(3.0, 1), (0.5, 0)], +) +@respx.mock +def test_flux2_pro_generation_bills_first_megapixel_per_returned_image( + output_mp: float, additional_output_megapixels: int +) -> None: + respx.post(url__regex=_flux2_route("flux-2-pro")).mock( + return_value=httpx.Response(200, json=_metered_body(output_mp, 0.25, num_images=2)) + ) + + response: Final = litellm.image_generation( + model="azure_ai/flux.2-pro", prompt="A red fox", api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response._hidden_params["response_cost"] == pytest.approx( + 2 + * ( + _price("azure_ai/flux.2-pro", "output_cost_per_first_megapixel") + + additional_output_megapixels * _price("azure_ai/flux.2-pro", "output_cost_per_megapixel") + ) + + _price("azure_ai/flux.2-pro", "input_cost_per_megapixel") + ) + + @pytest.mark.asyncio async def test_async_flux2_pro_generation_bills_the_reported_megapixels( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -349,9 +375,11 @@ def test_flux2_flex_generation_bills_every_reported_megapixel_at_one_rate() -> N [ ( {}, - lambda price: price("output_cost_per_first_megapixel") - + price("output_cost_per_megapixel") - + price("input_cost_per_megapixel"), + lambda price: ( + price("output_cost_per_first_megapixel") + + price("output_cost_per_megapixel") + + price("input_cost_per_megapixel") + ), ), ({"output_cost_per_image": 0.05}, lambda _: 0.05), ( @@ -359,7 +387,11 @@ def test_flux2_flex_generation_bills_every_reported_megapixel_at_one_rate() -> N lambda price: price("output_cost_per_first_megapixel") + 0.02 + price("input_cost_per_megapixel"), ), ( - {"output_cost_per_first_megapixel": 0.04, "output_cost_per_megapixel": 0.02, "input_cost_per_megapixel": 0.01}, + { + "output_cost_per_first_megapixel": 0.04, + "output_cost_per_megapixel": 0.02, + "input_cost_per_megapixel": 0.01, + }, lambda _: 0.04 + 0.02 + 0.01, ), ], From a8bfa64b2564d50c1c3da235a3126613b607ae4b Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 14:30:13 -0700 Subject: [PATCH 3/4] refactor(azure_ai): read FLUX.2 megapixel prices through a typed model --- .../image_generation/cost_calculator.py | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 5f421a7e2cd..92d9cb7a8ce 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -14,11 +14,6 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( from litellm.types.utils import ImageResponse, ModelInfo _MEGAPIXEL_USAGE_HIDDEN_PARAM: Final = "megapixel_usage" -_MEGAPIXEL_PRICE_KEYS: Final = ( - "output_cost_per_first_megapixel", - "output_cost_per_megapixel", - "input_cost_per_megapixel", -) class _MegapixelUsage(BaseModel): @@ -28,6 +23,14 @@ class _MegapixelUsage(BaseModel): output_mp: float = Field(gt=0) +class _MegapixelPrices(BaseModel): + model_config = ConfigDict(frozen=True) + + output_cost_per_megapixel: float + output_cost_per_first_megapixel: float | None = None + input_cost_per_megapixel: float | None = None + + class _MeteredImageResponse(BaseModel): model_config = ConfigDict(frozen=True) @@ -46,21 +49,19 @@ def record_megapixel_usage(response: object, response_body: bytes | Mapping[str, set_hidden_param(response, _MEGAPIXEL_USAGE_HIDDEN_PARAM, usage.model_dump()) -def _megapixel_cost(prices: ModelInfo, image_response: ImageResponse) -> float | None: - output_rate: Final = get_cost_per_unit(prices, "output_cost_per_megapixel", default_value=None) - if output_rate is None: - return None +def _megapixel_cost(model_info: ModelInfo, image_response: ImageResponse) -> float | None: try: + prices: Final = _MegapixelPrices.model_validate(model_info) usage: Final = _MegapixelUsage.model_validate(image_response.hidden_params.get(_MEGAPIXEL_USAGE_HIDDEN_PARAM)) except ValidationError: return None - first_rate: Final = get_cost_per_unit(prices, "output_cost_per_first_megapixel", default_value=None) - input_rate: Final = get_cost_per_unit(prices, "input_cost_per_megapixel") or 0.0 + first_rate: Final = prices.output_cost_per_first_megapixel images: Final = max(len(image_response.data or ()), 1) output_mp_per_image: Final = usage.output_mp / images return images * ( - (output_rate if first_rate is None else first_rate) + output_rate * (math.ceil(output_mp_per_image) - 1) - ) + input_rate * math.ceil(usage.input_mp or 0) + (prices.output_cost_per_megapixel if first_rate is None else first_rate) + + prices.output_cost_per_megapixel * (math.ceil(output_mp_per_image) - 1) + ) + (prices.input_cost_per_megapixel or 0.0) * math.ceil(usage.input_mp or 0) def _input_cost_per_pixel(resolved: ModelInfo) -> float: @@ -102,7 +103,7 @@ def cost_calculator( return token_based_cost bills_megapixels: Final = model_info is None or any( - model_info.get(key) is not None for key in _MEGAPIXEL_PRICE_KEYS + model_info.get(field) is not None for field in _MegapixelPrices.model_fields ) megapixel_cost: Final = _megapixel_cost(_model_info, image_response) if bills_megapixels else None if megapixel_cost is not None: From 723a9fe542cbc4819b6312565e7f70d89f2395ac Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Wed, 28 Oct 2026 21:49:11 -0700 Subject: [PATCH 4/4] fix(azure_ai): preserve fallback for empty FLUX.2 image responses --- .../azure_ai/image_generation/cost_calculator.py | 4 +++- .../test_azure_ai_flux2_image_generation.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 92d9cb7a8ce..b3d844f17cd 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -50,13 +50,15 @@ def record_megapixel_usage(response: object, response_body: bytes | Mapping[str, def _megapixel_cost(model_info: ModelInfo, image_response: ImageResponse) -> float | None: + images: Final = len(image_response.data or ()) + if not images: + return None try: prices: Final = _MegapixelPrices.model_validate(model_info) usage: Final = _MegapixelUsage.model_validate(image_response.hidden_params.get(_MEGAPIXEL_USAGE_HIDDEN_PARAM)) except ValidationError: return None first_rate: Final = prices.output_cost_per_first_megapixel - images: Final = max(len(image_response.data or ()), 1) output_mp_per_image: Final = usage.output_mp / images return images * ( (prices.output_cost_per_megapixel if first_rate is None else first_rate) diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 5a010f81556..4b85081f098 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -334,6 +334,20 @@ def test_flux2_pro_generation_bills_first_megapixel_per_returned_image( ) +@respx.mock +def test_flux2_pro_generation_without_returned_images_has_no_cost() -> None: + respx.post(url__regex=_flux2_route("flux-2-pro")).mock( + return_value=httpx.Response(200, json=_metered_body(3.0, 1.5, num_images=0)) + ) + + response: Final = litellm.image_generation( + model="azure_ai/flux.2-pro", prompt="A red fox", n=4, api_base=FOUNDRY_BASE, api_key="test-key" + ) + + assert response.data == [] + assert response._hidden_params["response_cost"] == 0.0 + + @pytest.mark.asyncio async def test_async_flux2_pro_generation_bills_the_reported_megapixels( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch