From 61a361abbf1ace8d0e553a680a0cc906b310818a Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Fri, 25 Sep 2026 01:42:16 +0000 Subject: [PATCH] fix(images): keep deployment pixel prices on image edit logging so proxy edits bill the deployment rate --- litellm/images/main.py | 5 ++ .../pricing/test_image_pixel_pricing.py | 49 +++++++++++++++++++ ...test_azure_ai_image_edit_transformation.py | 40 +++++++++++++++ 3 files changed, 94 insertions(+) create mode 100644 tests/integration/pricing/test_image_pixel_pricing.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 1f722eb752a..f16a5c22187 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -51,6 +51,7 @@ from litellm.types.llms.openai import ImageGenerationRequestQuality from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( LITELLM_IMAGE_VARIATION_PROVIDERS, + CustomPricingLiteLLMParams, LlmProviders, all_litellm_params, ) @@ -883,6 +884,10 @@ def image_edit( optional_params=dict(image_edit_request_params), litellm_params={ **image_edit_request_params, + **litellm_params.model_dump( + include=frozenset(CustomPricingLiteLLMParams.model_fields), + exclude_none=True, + ), "litellm_call_id": litellm_call_id, "model_info": model_info, }, diff --git a/tests/integration/pricing/test_image_pixel_pricing.py b/tests/integration/pricing/test_image_pixel_pricing.py new file mode 100644 index 00000000000..f632cb8714c --- /dev/null +++ b/tests/integration/pricing/test_image_pixel_pricing.py @@ -0,0 +1,49 @@ +import json +from typing import Final + +import pytest + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.database import read_rows +from tests.integration._support.wire import Reply, Request, wire_server + +_PIXELS_PER_MEGAPIXEL: Final = 1024 * 1024 + + +@pytest.mark.parametrize("input_cost_per_pixel", (1.5e-07, 0.0)) +def test_flux2_edit_bills_azure_reported_megapixels_at_the_deployment_rate( + gateway: Gateway, input_cost_per_pixel: float +) -> None: + def respond(request: Request) -> Reply: + assert request.method == "POST" + return Reply( + body=json.dumps( + {"data": [{"b64_json": "aW1n"}], "request_meta": {"input_mp": 0.39, "output_mp": 1.0}} + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/FLUX.2-flex", + api_base=wire.url, + api_key="synthetic-azure-key", + input_cost_per_pixel=input_cost_per_pixel, + ) + response: Final = gateway.request_multipart( + "/v1/images/edits", + {"model": model, "prompt": "add a hat", "size": "1024x1024"}, + {"image": ("reference.png", b"reference", "image/png")}, + ) + assert response.status_code == 200, response.text + expected: Final = input_cost_per_pixel * _PIXELS_PER_MEGAPIXEL * (0.39 + 1.0) + if expected: + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected), response.text + else: + assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0"), response.text + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(rows[0]["spend"]) == pytest.approx(expected), response.text 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 a74d93846e7..0ec081d9864 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 @@ -5,6 +5,7 @@ from typing import Final import httpx import pytest +import respx import litellm from litellm.images.utils import ImageEditRequestUtils @@ -215,6 +216,45 @@ def test_flux2_image_edit_bills_the_megapixels_azure_reports(input_mp: float, ou assert response._hidden_params["response_cost"] == pytest.approx(rate * 1024 * 1024 * (input_mp + output_mp)) +@pytest.mark.parametrize("input_cost_per_pixel", (1.5e-07, 0.0)) +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.asyncio +async def test_flux2_image_edit_bills_azure_megapixels_at_the_deployment_rate( + input_cost_per_pixel: float, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary. + litellm, "disable_aiohttp_transport", True + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(url__startswith="https://example.services.ai.azure.com/").respond( + json={"data": [{"b64_json": "aW1n"}], "request_meta": {"input_mp": 0.39, "output_mp": 1.0}} + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "flux-deployment", + "litellm_params": { + "model": "azure_ai/FLUX.2-flex", + "api_key": "test-key", + "api_base": "https://example.services.ai.azure.com", + "input_cost_per_pixel": input_cost_per_pixel, + }, + } + ] + ) + + response: Final = await router.aimage_edit( + model="flux-deployment", + image=b"reference", + prompt="Add a hat", + size="1024x1024", + ) + + assert response._hidden_params["response_cost"] == pytest.approx(input_cost_per_pixel * 1024 * 1024 * 1.39) + + def test_flux2_image_edit_accepts_and_drops_openai_only_parameters(): optional_params: Final = ImageEditRequestUtils.get_optional_params_image_edit( model="FLUX.2-pro",