diff --git a/litellm/llms/xai/image_generation/cost_calculator.py b/litellm/llms/xai/image_generation/cost_calculator.py index c102224bfac..5cf575ae743 100644 --- a/litellm/llms/xai/image_generation/cost_calculator.py +++ b/litellm/llms/xai/image_generation/cost_calculator.py @@ -28,6 +28,9 @@ def cost_calculator( call_type: str | None = None, model_info: ModelInfo | None = None, ) -> float | None: + if model_info is not None and model_info.get("input_cost_per_pixel") is not None: + return None + caller_set_prices: Final = model_info is not None and any( model_info.get(key) is not None for key in ( diff --git a/tests/unit/llms/xai/image_generation/test_cost_calculator.py b/tests/unit/llms/xai/image_generation/test_cost_calculator.py index 83de6b9b9a2..70246534368 100644 --- a/tests/unit/llms/xai/image_generation/test_cost_calculator.py +++ b/tests/unit/llms/xai/image_generation/test_cost_calculator.py @@ -1,3 +1,4 @@ +from datetime import datetime from typing import Final import httpx @@ -6,6 +7,7 @@ from openai import AsyncOpenAI, OpenAI import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object from litellm.types.utils import ImageResponse @@ -141,6 +143,92 @@ def test_deployment_image_price_overrides_billed_ticks(price_key: str, resolutio ) == pytest.approx(0.83 * 3) +@pytest.mark.parametrize("ticks", (None, -1, True, "600000000", 0, 600000000)) +@pytest.mark.parametrize("pixel_rate", (0.0, 0.000001)) +@pytest.mark.parametrize( + "size,pixels,n", ((None, 1024 * 1024, 3), ("1024x1024", 1024 * 1024, 1), ("512x256", 512 * 256, 2)) +) +def test_deployment_pixel_price_overrides_ticks_and_resolution_prices( + ticks: object, pixel_rate: float, size: str | None, pixels: int, n: int +) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="grok-imagine-image-2.0", + custom_llm_provider="xai", + completion_response=_response(ticks), + size=size, + quality="low", + n=n, + model_info={"input_cost_per_pixel": pixel_rate, "output_cost_per_image_2K": 0.29}, + optional_params={"resolution": "2k"}, + call_type="image_generation", + ) == pytest.approx(pixel_rate * pixels * n) + + +@pytest.mark.parametrize("ticks", (None, 600000000)) +@pytest.mark.parametrize("price_key", ("input_cost_per_image", "output_cost_per_image")) +def test_deployment_flat_price_still_precedes_pixel_price(ticks: object, price_key: str) -> None: + assert CostCalculatorUtils.route_image_generation_cost_calculator( + model="unmapped-image-fixture", + custom_llm_provider="xai", + completion_response=_response(ticks), + n=2, + quality="low", + model_info={price_key: 0.83, "input_cost_per_pixel": 0.000001, "output_cost_per_image_2K": 0.29}, + optional_params={"resolution": "2k"}, + ) == pytest.approx(0.83 * 2) + + +@pytest.mark.asyncio +async def test_router_preserves_deployment_pixel_price_through_image_normalization() -> None: + ticks: Final = 600000000 + transport: Final = httpx.MockTransport( + lambda request: httpx.Response( + 200, + json={ + "created": 123, + "data": [{"url": "https://image.test/output.png"}], + "usage": {"cost_in_usd_ticks": ticks}, + }, + ) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "pixel-fixture", + "litellm_params": { + "model": "xai/grok-imagine-image-2.0", + "api_key": "test", + "api_base": "https://image.test/v1", + }, + "model_info": {"input_cost_per_pixel": 0.000001}, + } + ] + ) + logging_obj: Final = Logging( + model="grok-imagine-image-2.0", + messages=[], + stream=False, + call_type="aimage_generation", + start_time=datetime(2026, 1, 1), + litellm_call_id="pixel-override", + function_id="pixel-override", + ) + async with AsyncOpenAI( + api_key="test", base_url="https://image.test/v1", http_client=httpx.AsyncClient(transport=transport) + ) as client: + response: Final = await router.aimage_generation( + model="pixel-fixture", + prompt="circle", + quality="low", + n=1, + extra_body={"resolution": "2k"}, + client=client, + litellm_logging_obj=logging_obj, + ) + assert getattr(response.usage, "cost_in_usd_ticks") == ticks + assert logging_obj.response_cost_calculator(result=response) == pytest.approx(1024 * 1024 * 0.000001) + + def test_shared_resolution_prices_are_used_when_ticks_are_absent() -> None: pricing: Final = litellm.get_model_info("xai/grok-imagine-image-2.0") assert CostCalculatorUtils.route_image_generation_cost_calculator(