fix(azure_ai): bill FLUX.2 Pro edit references on the Pro ref meter

This commit is contained in:
shrey kharbanda 2026-09-24 23:38:53 +00:00
parent 9505cb1553
commit ee0fd20095
4 changed files with 40 additions and 0 deletions

View file

@ -10954,6 +10954,7 @@
]
},
"azure_ai/flux.2-pro": {
"input_cost_per_pixel": 1.430511474609375e-08,
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.04,

View file

@ -10954,6 +10954,7 @@
]
},
"azure_ai/flux.2-pro": {
"input_cost_per_pixel": 1.430511474609375e-08,
"litellm_provider": "azure_ai",
"mode": "image_generation",
"output_cost_per_image": 0.04,

View file

@ -1,5 +1,6 @@
import base64
import io
import tempfile
from typing import Final
from litellm.llms.azure_ai.image_edit.flux2_transformation import AzureFoundryFlux2ImageEditConfig
@ -37,3 +38,22 @@ def test_convert_image_to_base64_rewinds_seekable_streams_before_encoding() -> N
encoded: Final = AzureFoundryFlux2ImageEditConfig()._convert_image_to_base64(stream)
assert base64.b64decode(encoded) == b"png-bytes"
def test_convert_image_to_base64_sends_the_full_image_on_every_retry() -> None:
config: Final = AzureFoundryFlux2ImageEditConfig()
stream: Final = io.BytesIO(b"png-bytes")
first: Final = config._convert_image_to_base64(stream)
retry: Final = config._convert_image_to_base64(stream)
assert base64.b64decode(retry) == base64.b64decode(first) == b"png-bytes"
def test_convert_image_to_base64_accepts_named_temporary_files() -> None:
with tempfile.NamedTemporaryFile() as upload:
upload.write(b"png-bytes")
encoded: Final = AzureFoundryFlux2ImageEditConfig()._convert_image_to_base64(upload)
assert base64.b64decode(encoded) == b"png-bytes"

View file

@ -520,3 +520,21 @@ def test_flux2_cost_optional_params_dimensions_beat_response_size() -> None:
def test_flux2_cost_rejects_non_image_response() -> None:
with pytest.raises(TypeError, match="must be of type ImageResponse"):
cost_calculator(model="FLUX.2-flex", image_response=object())
def test_flux2_pro_edit_bills_references_on_the_pro_ref_meter() -> None:
pro_pixel_rate: Final = litellm.model_cost["azure_ai/flux.2-pro"]["input_cost_per_pixel"]
assert pro_pixel_rate * _ONE_MEGAPIXEL == pytest.approx(0.015), (
"Azure Retail Prices API, product 'Azure BFL Flux Models', meter 'Flux 2 Ref MP Megapixel' is $0.015 "
"per MP, checked 2026-09-24"
)
cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
model="flux.2-pro",
completion_response=_edit_response(),
custom_llm_provider="azure_ai",
size="1024x1024",
call_type="image_edit",
)
assert cost == pytest.approx(litellm.model_cost["azure_ai/flux.2-pro"]["output_cost_per_image"] + 0.015 * 2)