fix(images): keep deployment pixel prices on image edit logging so proxy edits bill the deployment rate

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,
             },

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

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",