fix(azure-ai): price FLUX.2 Flex by its resolved cost map key

This commit is contained in:
mateo-berri 2026-09-18 11:42:52 -07:00
parent 4595b4f62f
commit cf08440557
2 changed files with 14 additions and 4 deletions

View file

@ -42,9 +42,6 @@ def cost_calculator(
if input_cost_per_pixel:
from litellm.cost_calculator import default_image_cost_calculator
cost_model: Final = (
model if model.startswith(f"{litellm.LlmProviders.AZURE_AI.value}/") else f"azure_ai/{model}"
)
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 = (
@ -53,7 +50,7 @@ def cost_calculator(
else size or image_response.size
)
return default_image_cost_calculator(
model=cost_model,
model=_model_info["key"],
custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
size=pixel_size,
n=num_images,

View file

@ -172,6 +172,19 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi
) == pytest.approx(5e-08 * 2048 * 1024 * 2)
def test_flux2_flex_cost_accepts_lowercase_model_spelling():
response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")])
cost: Final = litellm.completion_cost(
model="azure_ai/flux.2-flex",
completion_response=response,
optional_params={"width": 1536, "height": 1024, "num_images": 2},
call_type="image_generation",
)
assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2)
def test_flux2_response_preserves_mapped_dimensions():
config = AzureFoundryFluxImageGenerationConfig()
params = config.map_openai_params(